Update app.py

app.py

@@ -31,6 +31,7 @@ import diffusers
 # init
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print(device)
 base_model = "black-forest-labs/FLUX.1-dev"
 
 # load pipe
@@ -116,12 +117,16 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
     gr.Info("Starting process")
     img2img_model = False
     orginal_image = None
-    if image_url:
-        orginal_image = load_image(image_url)
+    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"):
@@ -130,7 +135,7 @@ 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()
+        # img2img_pipe.unload_lora_weights()
         txt2img_pipe.unload_lora_weights()
     
     lora_configs = None
@@ -162,7 +167,6 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
                         try:
                             if img2img_model:
                                 img2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True,  adapter_name=lora_name)
-                                img2img_pipe.set
                             else:
                                 txt2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True,  adapter_name=lora_name)
                         except:
@@ -187,6 +191,7 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
             joint_attention_kwargs = {"scale": 1}
             
             if orginal_image:
+                img2img_pipe.to(device)
                 final_image = img2img_pipe(
                     prompt=prompt,
                     image=orginal_image,
@@ -199,6 +204,7 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
                     joint_attention_kwargs=joint_attention_kwargs
                 ).images[0]
             else:
+                txt2img_pipe.to(device)
                 final_image = txt2img_pipe(
                     prompt=prompt,
                     num_inference_steps=steps,