Update app.py

app.py

@@ -21,7 +21,7 @@ THEME_PRESETS = {
     "orange-red": {"primary": "#c2410c", "secondary": "#ea580c"},
     "violet":     {"primary": "#7c3aed", "secondary": "#8b5cf6"},
     "ocean":      {"primary": "#0369a1", "secondary": "#0ea5e9"},
-    "emerald":    {"primary": "#047857", "secondary": "#10b981"},
+    "emerald":    {"primary": "#226105", "secondary": "#73ff00"},
     "rose":       {"primary": "#be185d", "secondary": "#ec4899"},
     "amber":      {"primary": "#b45309", "secondary": "#f59e0b"},
     "cyan":       {"primary": "#0e7490", "secondary": "#06b6d4"},
@@ -289,34 +289,35 @@ def generate(
 
     # ZeroGPU: LoRA loading must happen inside GPU scope
     pipe.to("cuda")
-    pipe.unload_lora_weights()
-
-    if speed_mode == "Fast · 8 steps":
-        pipe.load_lora_weights(LIGHTNING_REPO, weight_name=LIGHTNING_WEIGHT, adapter_name="lightning")
-        pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
-        pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
-    else:
-        pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
-        pipe.set_adapters(["style"], adapter_weights=[lora_scale])
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    w, h = aspect_to_wh(aspect)
-
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    image = pipe(
-        prompt=prompt_in,
-        negative_prompt=neg_prompt,
-        num_inference_steps=steps,
-        true_cfg_scale=cfg,
-        width=w,
-        height=h,
-        generator=generator,
-    ).images[0]
-
-    pipe.to("cpu")
-    torch.cuda.empty_cache()
+    try:
+        pipe.unload_lora_weights()
+
+        if speed_mode == "Fast · 8 steps":
+            pipe.load_lora_weights(LIGHTNING_REPO, weight_name=LIGHTNING_WEIGHT, adapter_name="lightning")
+            pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
+            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
+        else:
+            pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
+            pipe.set_adapters(["style"], adapter_weights=[lora_scale])
+
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+
+        w, h = aspect_to_wh(aspect)
+
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        image = pipe(
+            prompt=prompt_in,
+            negative_prompt=neg_prompt,
+            num_inference_steps=steps,
+            true_cfg_scale=cfg,
+            width=w,
+            height=h,
+            generator=generator,
+        ).images[0]
+    finally:
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
 
     # update history
     history_state = history_state or []
@@ -397,13 +398,16 @@ def add_custom_lora(custom_text: str):
             </div>
           </div>
         </div>"""
-        return gr.update(visible=True, value=card_html), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {weight}", existing
+        new_gallery = [(l["image"], l["title"]) for l in loras]
+        return (gr.update(visible=True, value=card_html), gr.update(visible=True), 
+                gr.update(value=new_gallery, selected_index=None), f"Custom: {weight}", existing)
     except Exception as e:
         gr.Warning(str(e))
         return gr.update(visible=True, value=f"<span style='color:#e24b4a'>{e}</span>"), gr.update(visible=True), gr.update(), "", None
 
 def remove_custom_lora():
-    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None
+    gallery_reset = [(l["image"], l["title"]) for l in loras]
+    return gr.update(visible=False), gr.update(visible=False), gr.update(value=gallery_reset), "", None
 
 generate.zerogpu = True