fix pipeline load

app.py

@@ -26,13 +26,13 @@ weight_dtype = torch.bfloat16
 args = OmegaConf.load('configs/omnitry_v1_unified.yaml')
 
 # init model
-transformer = FluxTransformer2DModel.from_pretrained('./FLUX.1-Fill-dev/transformer').requires_grad_(False).to(device, dtype=weight_dtype)
-vae = diffusers.AutoencoderKL.from_pretrained('./FLUX.1-Fill-dev/vae').requires_grad_(False).to(device, dtype=weight_dtype)
-text_encoder = transformers.CLIPTextModel.from_pretrained('./FLUX.1-Fill-dev/text_encoder').requires_grad_(False).to(device, dtype=weight_dtype)
-text_encoder_2 = transformers.T5EncoderModel.from_pretrained('./FLUX.1-Fill-dev/text_encoder_2').requires_grad_(False).to(device, dtype=weight_dtype)
-scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained('./FLUX.1-Fill-dev/scheduler')
-tokenizer = transformers.CLIPTokenizer.from_pretrained('./FLUX.1-Fill-dev/tokenizer')
-tokenizer_2 = transformers.T5TokenizerFast.from_pretrained('./FLUX.1-Fill-dev/tokenizer_2')
+transformer = FluxTransformer2DModel.from_pretrained('black-forest-labs/FLUX.1-Fill-dev', subfolder='transformer').requires_grad_(False).to(device, dtype=weight_dtype)
+pipeline = FluxFillPipeline.from_pretrained(
+    'black-forest-labs/FLUX.1-Fill-dev',
+    transformer=transformer, 
+    torch_dtype=weight_dtype
+).to(device)
+
 
 # insert LoRA
 lora_config = LoraConfig(
@@ -81,17 +81,6 @@ for n, m in transformer.named_modules():
     if isinstance(m, peft.tuners.lora.layer.Linear):
         m.forward = create_hacked_forward(m)
 
-# init pipeline
-pipeline = FluxFillPipeline(
-    transformer=transformer.eval(), 
-    scheduler=copy.deepcopy(scheduler),
-    vae=vae,
-    text_encoder=text_encoder,
-    text_encoder_2=text_encoder_2,
-    tokenizer=tokenizer,
-    tokenizer_2=tokenizer_2,
-)
-
 
 def seed_everything(seed=0):
     random.seed(seed)