Update app.py

app.py

@@ -40,6 +40,7 @@ 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)
+txt2img_pipe.to(device)
 
 
 MAX_SEED = 2**32 - 1
@@ -123,6 +124,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"):
+        txt2img_pipe.to(device)
         txt2img_pipe.unload_lora_weights()
     print(device)
     lora_configs = None
@@ -139,7 +141,7 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
            
             with calculateDuration("Loading LoRA weights"):
                 adapter_weights = []
-
+                txt2img_pipe.to(device)
                 for idx, lora_info in enumerate(lora_configs):
                     lora_repo = lora_info.get("repo")
                     weights = lora_info.get("weights")
@@ -165,10 +167,12 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
     try:
         gr.Info("Start to generate images ...")
         with calculateDuration(f"Make a new generator: {seed}"):
+            txt2img_pipe.to(device)
             generator = torch.Generator(device=device).manual_seed(seed)
         print(device)
         with calculateDuration("Generating image"):
             # Generate image
+            txt2img_pipe.to(device)
             joint_attention_kwargs = {"scale": 1}
             final_image = txt2img_pipe(
                 prompt=prompt,