Update app.py

app.py

@@ -41,10 +41,6 @@ good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtyp
 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
 
@@ -118,15 +114,7 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
     img2img_model = False
     orginal_image = None
     print(device)
-    if image_url and image_url != "":
-        print("img2img")
-        orginal_image = load_image(image_url).to(device)
-        img2img_model = True
-        img2img_pipe.to(device)
-    else:
-        print("txt2img")
-        txt2img_pipe.to(device)
-        
+  
     # Set random seed for reproducibility
     if randomize_seed:
         with calculateDuration("Set random seed"):
@@ -135,9 +123,8 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
     # Load LoRA weights
     gr.Info("Start to load LoRA ...")
     with calculateDuration("Unloading LoRA"):
-        # img2img_pipe.unload_lora_weights()
         txt2img_pipe.unload_lora_weights()
-    
+    print(device)
     lora_configs = None
     adapter_names = []
     lora_names = []
@@ -165,19 +152,13 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
                     adapter_weights.append(adapter_weight)
                     if lora_repo and weights and adapter_name:
                         try:
-                            if img2img_model:
-                                img2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True,  adapter_name=lora_name)
-                            else:
-                                txt2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True,  adapter_name=lora_name)
+                            txt2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True,  adapter_name=lora_name)
                         except:
                             print("load lora error")
                 
                 # set lora weights
                 if len(lora_names) > 0:
-                    if img2img_model:
-                        img2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
-                    else:
-                        txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
+                    txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
     
     # Generate image
     error_message = ""
@@ -185,36 +166,20 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
         gr.Info("Start to generate images ...")
         with calculateDuration(f"Make a new generator: {seed}"):
             generator = torch.Generator(device=device).manual_seed(seed)
-        
+        print(device)
         with calculateDuration("Generating image"):
             # Generate image
             joint_attention_kwargs = {"scale": 1}
-            
-            if orginal_image:
-                img2img_pipe.to(device)
-                final_image = img2img_pipe(
-                    prompt=prompt,
-                    image=orginal_image,
-                    strength=image_strength,
-                    num_inference_steps=steps,
-                    guidance_scale=cfg_scale,
-                    width=width,
-                    height=height,
-                    generator=generator,
-                    joint_attention_kwargs=joint_attention_kwargs
-                ).images[0]
-            else:
-                txt2img_pipe.to(device)
-                final_image = txt2img_pipe(
-                    prompt=prompt,
-                    num_inference_steps=steps,
-                    guidance_scale=cfg_scale,
-                    width=width,
-                    height=height,
-                    max_sequence_length=512,
-                    generator=generator,
-                    joint_attention_kwargs=joint_attention_kwargs
-                ).images[0]
+            final_image = txt2img_pipe(
+                prompt=prompt,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                max_sequence_length=512,
+                generator=generator,
+                joint_attention_kwargs=joint_attention_kwargs
+            ).images[0]
  
     except Exception as e:
         error_message =  str(e)