Update app.py

app.py

@@ -9,7 +9,7 @@ import logging
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
 from huggingface_hub import login
 from diffusers.utils import load_image
-
+from lora_loading_patch import load_lora_into_transformer
 import time
 from datetime import datetime
 from io import BytesIO
@@ -34,19 +34,13 @@ base_model = "black-forest-labs/FLUX.1-dev"
 # load pipe
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
-# img2img model
-img2img = AutoPipelineForImage2Image.from_pretrained(base_model,
-                                                      vae=good_vae,
-                                                      transformer=pipe.transformer,
-                                                      text_encoder=pipe.text_encoder,
-                                                      tokenizer=pipe.tokenizer,
-                                                      text_encoder_2=pipe.text_encoder_2,
-                                                      tokenizer_2=pipe.tokenizer_2,
-                                                      torch_dtype=dtype
-                                                     )
+txt2img_pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
+txt2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
 
+# img2img model
+img2img_pipe = AutoPipelineForImage2Image.from_pretrained(base_model, vae=good_vae, transformer=txt2img_pipe.transformer, text_encoder=txt2img_pipe.text_encoder, tokenizer=txt2img_pipe.tokenizer, text_encoder_2=txt2img_pipe.text_encoder_2, tokenizer_2=txt2img_pipe.tokenizer_2, torch_dtype=dtype)
+img2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
 
 
 MAX_SEED = 2**32 - 1
@@ -76,15 +70,19 @@ def generate_image(orginal_image, prompt, adapter_names,  steps, seed, image_str
 
     
     gr.Info("Start to generate images ...")
-    with calculateDuration(f"Make a new generator:{seed}"):
-        pipe.to(device)
+    with calculateDuration(f"Make a new generator: {seed}"):
+        if orginal_image:
+            img2img_pipe.to(device)
+        else:
+            txt2img_pipe.to(device)
+        
         generator = torch.Generator(device=device).manual_seed(seed)
     
     with calculateDuration("Generating image"):
         # Generate image
-        joint_attention_kwargs = {"scale": 1}    
+        joint_attention_kwargs = {"scale": 1}
         if orginal_image:
-            generated_image = img2img(
+            generated_image = img2img_pipe(
                 prompt=prompt,
                 image=orginal_image,
                 strength=image_strength,
@@ -96,7 +94,7 @@ def generate_image(orginal_image, prompt, adapter_names,  steps, seed, image_str
                 joint_attention_kwargs=joint_attention_kwargs
             ).images[0]
         else:
-            generated_image = pipe(
+            generated_image = txt2img_pipe(
                 prompt=prompt,
                 num_inference_steps=steps,
                 guidance_scale=cfg_scale,
@@ -189,18 +187,18 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
                     if lora_repo and weights and adapter_name:
                         try:
                             if img2img_model:
-                                img2img.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                                img2img_pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                             else:
-                                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                                txt2img_pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                         except:
                             print("load lora error")
                 
                 # set lora weights
                 if len(adapter_names) > 0:
                     if img2img_model:
-                        img2img.set_adapters(adapter_names, adapter_weights=adapter_weights)
+                        img2img_pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
                     else:
-                        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+                        txt2img_pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
 
     
     # Generate image