attempt to load lora on cpu

app.py

@@ -49,6 +49,20 @@ def update_selection(evt: gr.SelectData):
     )
 
 @spaces.GPU(duration=80)
+def generate_image(prompt, trigger_word, steps, cfg_scale, width, height, generator, lora_scale, progress):
+    with calculateDuration("Generating image"):
+        # Generate image
+        image = pipe(
+            prompt=f"{prompt} {trigger_word}",
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
+    return image
+
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
@@ -69,24 +83,11 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         generator = torch.Generator(device="cuda").manual_seed(seed)
-
-    with calculateDuration("Generating image"):
-        # Generate image
-        image = pipe(
-            prompt=f"{prompt} {trigger_word}",
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
-
-    yield image, seed
-    
-    with calculateDuration("Unloading weights"):
-        pipe.unload_lora_weights()
     
+    image = generate_image(prompt, trigger_word, steps, cfg_scale, width, height, generator, lora_scale, progress)
+
+    pipe.unload_lora_weights()
+    return image, seed  
 
 css = '''
 #gen_btn{height: 100%}