update model

app.py

@@ -6,6 +6,8 @@ from kontext.pipeline_flux_kontext import FluxKontextPipeline
 from kontext.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from diffusers import FluxTransformer2DModel
 import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 def resize_by_bucket(images_pil, resolution=512):
     assert len(images_pil) > 0, "images_pil 不能为空"
@@ -49,8 +51,9 @@ flux_pipeline.vae.to(device).to(torch.bfloat16)
 flux_pipeline.text_encoder.to(device).to(torch.bfloat16)
 flux_pipeline.text_encoder_2.to(device).to(torch.bfloat16)
 flux_pipeline.scheduler.config.stochastic_sampling = False
-finetuned_path = "NoobDoge/Multi_Ref"
-flux_pipeline.transformer = FluxTransformer2DModel.from_pretrained(finetuned_path,subfolder='transformer', torch_dtype=torch.bfloat16)
+ckpt_path = hf_hub_download("NoobDoge/Multi_Ref_Model", "full_ema_model.safetensors")
+new_weight = load_file(ckpt_path)
+flux_pipeline.transformer.load_state_dict(new_weight)
 flux_pipeline.transformer.to(device).to(torch.bfloat16)
 
 MAX_SEED = np.iinfo(np.int32).max

requirements.txt

@@ -4,4 +4,5 @@ invisible_watermark
 torch
 transformers
 xformers
-sentencepiece
+sentencepiece
+safetensors