feat: Quick Edit tab, LoRA Dataset Generator (50-150 captioned poses), fix LoRA reload for ZeroGPU

app.py

@@ -157,8 +157,7 @@ def infer(
         if not prompt or not prompt.strip():
             raise gr.Error("Enter a prompt!")
 
-        # LoRA activation
-        global LOADED_ADAPTERS
+        # LoRA activation — always reload to survive ZeroGPU tensor packing
         active_adapters = []
         active_weights = []
 
@@ -166,28 +165,31 @@ def infer(
         style = LORA_MAP.get(lora_name)
         if style and style["adapter_name"]:
             aname = style["adapter_name"]
-            if aname not in LOADED_ADAPTERS:
-                print(f"Loading LoRA: {style['title']}")
+            try:
+                print(f"Loading LoRA: {style['title']} (strength={lora_strength})")
                 pipe.load_lora_weights(style["repo"], weight_name=style["weights"], adapter_name=aname)
-                LOADED_ADAPTERS.add(aname)
+            except ValueError:
+                pass  # already loaded
             active_adapters.append(aname)
             active_weights.append(lora_strength)
 
         # Custom LoRA
         if custom_lora_repo and custom_lora_repo.strip() and custom_lora_file and custom_lora_file.strip():
             cname = "custom-lora"
-            if cname not in LOADED_ADAPTERS:
+            try:
                 print(f"Loading custom LoRA: {custom_lora_repo}/{custom_lora_file}")
                 pipe.load_lora_weights(custom_lora_repo.strip(), weight_name=custom_lora_file.strip(), adapter_name=cname)
-                LOADED_ADAPTERS.add(cname)
+            except ValueError:
+                pass  # already loaded
             active_adapters.append(cname)
             active_weights.append(custom_lora_strength)
 
         if active_adapters:
             pipe.set_adapters(active_adapters, adapter_weights=active_weights)
-            print(f"Active: {list(zip(active_adapters, active_weights))}")
+            print(f"Active LoRAs: {list(zip(active_adapters, active_weights))}")
         else:
             pipe.disable_lora()
+            print("No LoRA active")
 
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
@@ -230,14 +232,14 @@ def generate_character_sheet(
         ref = pil_images[0]
         prefix = custom_prefix.strip() + " " if custom_prefix and custom_prefix.strip() else ""
 
-        # Activate LoRA
-        global LOADED_ADAPTERS
+        # Activate LoRA — always reload for ZeroGPU
         style = LORA_MAP.get(lora_name)
         if style and style["adapter_name"]:
             aname = style["adapter_name"]
-            if aname not in LOADED_ADAPTERS:
+            try:
                 pipe.load_lora_weights(style["repo"], weight_name=style["weights"], adapter_name=aname)
-                LOADED_ADAPTERS.add(aname)
+            except ValueError:
+                pass
             pipe.set_adapters([aname], adapter_weights=[lora_strength])
         else:
             pipe.disable_lora()
@@ -325,6 +327,55 @@ with gr.Blocks(css=css) as demo:
                     outputs=[output_image, seed_output],
                 )
 
+            # ==================== SIMPLE IMAGE EDIT TAB ====================
+            with gr.TabItem("Quick Image Edit"):
+                gr.Markdown("Upload one image and describe what to change. No LoRA — pure Klein 9B editing.")
+                with gr.Row():
+                    with gr.Column():
+                        qe_image = gr.Image(label="Source Image", type="pil", sources=["upload"], height=350)
+                        qe_prompt = gr.Textbox(label="Edit instruction", lines=2,
+                                               placeholder="e.g. remove the shirt, change hair to blonde, add sunglasses...")
+                        with gr.Row():
+                            qe_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                            qe_guidance = gr.Slider(0.0, 10.0, value=1.0, step=0.1, label="Guidance")
+                            qe_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
+                        qe_rand = gr.Checkbox(value=True, label="Randomize seed")
+                        qe_btn = gr.Button("Edit Image", variant="primary", size="lg")
+                    with gr.Column():
+                        qe_output = gr.Image(label="Result", interactive=False, format="png", height=500)
+                        qe_seed_out = gr.Number(label="Seed Used", interactive=False)
+
+                @spaces.GPU
+                def quick_edit(image, prompt, steps, guidance, seed, randomize):
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                    try:
+                        if image is None:
+                            raise gr.Error("Upload an image!")
+                        if not prompt or not prompt.strip():
+                            raise gr.Error("Enter an edit instruction!")
+                        pipe.disable_lora()
+                        if randomize:
+                            seed = random.randint(0, MAX_SEED)
+                        w, h = update_dimensions(image)
+                        img = image.resize((w, h), Image.LANCZOS).convert("RGB")
+                        result = pipe(
+                            image=img, prompt=prompt,
+                            guidance_scale=guidance, width=w, height=h,
+                            num_inference_steps=steps,
+                            generator=torch.Generator(device=device).manual_seed(seed),
+                        ).images[0]
+                        return result, seed
+                    finally:
+                        gc.collect()
+                        torch.cuda.empty_cache()
+
+                qe_btn.click(
+                    fn=quick_edit,
+                    inputs=[qe_image, qe_prompt, qe_steps, qe_guidance, qe_seed, qe_rand],
+                    outputs=[qe_output, qe_seed_out],
+                )
+
             # ==================== CHARACTER SHEET TAB ====================
             with gr.TabItem("360 Character Sheet"):
                 gr.Markdown("Generate a multi-angle character turnaround from a single reference image. Produces 7 views: front/left/right face + front/left/right/back body.")
@@ -361,5 +412,150 @@ with gr.Blocks(css=css) as demo:
                     outputs=[cs_gallery],
                 )
 
+            # ==================== DATASET GENERATOR TAB ====================
+            with gr.TabItem("LoRA Dataset Generator"):
+                gr.Markdown("Generate a captioned image dataset from a single reference for LoRA training. Each image gets a text caption file.")
+
+                POSE_LIBRARY = [
+                    "standing facing the camera, neutral pose, arms at sides",
+                    "standing with arms crossed, confident pose",
+                    "standing with hands on hips",
+                    "standing, slight lean to the left, relaxed",
+                    "standing three-quarter view from the left",
+                    "standing three-quarter view from the right",
+                    "standing side profile, looking right",
+                    "standing side profile, looking left",
+                    "standing from behind, back view",
+                    "standing over the shoulder look, glancing back at camera",
+                    "sitting on a chair, legs crossed, relaxed",
+                    "sitting on the floor, legs extended",
+                    "sitting cross-legged on the ground",
+                    "sitting on a stool, leaning forward slightly",
+                    "sitting sideways on a chair, arm draped over backrest",
+                    "kneeling on one knee",
+                    "kneeling on both knees, upright posture",
+                    "leaning against a wall, arms crossed",
+                    "leaning against a wall, one foot up",
+                    "leaning forward with hands on knees",
+                    "walking towards the camera, mid-stride",
+                    "walking away from camera, back view mid-stride",
+                    "walking side view, profile mid-stride",
+                    "running towards the camera, dynamic pose",
+                    "looking up at the sky, chin raised",
+                    "looking down, contemplative",
+                    "head tilted to the left, slight smile",
+                    "head tilted to the right, serious expression",
+                    "laughing naturally, candid expression",
+                    "hands behind head, stretching",
+                    "one hand touching hair, casual pose",
+                    "hands in pockets, casual standing",
+                    "waving at camera, friendly gesture",
+                    "pointing at camera, direct gesture",
+                    "arms raised above head, celebratory",
+                    "crouching down, low angle",
+                    "bending forward slightly, looking at camera",
+                    "twisting torso, looking over shoulder",
+                    "dancing pose, one leg lifted",
+                    "yoga tree pose, balanced on one leg",
+                    "lying on back, looking up at camera from above",
+                    "lying on side, propped on elbow",
+                    "lying on stomach, chin in hands",
+                    "close-up portrait, direct eye contact",
+                    "close-up portrait, eyes looking away",
+                    "close-up portrait, slight smile",
+                    "close-up portrait, serious expression",
+                    "medium shot from waist up, arms at sides",
+                    "medium shot from waist up, one hand raised",
+                    "full body shot, standing tall, power pose",
+                ]
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        ds_ref = gr.Gallery(label="Reference Image", type="filepath", columns=1, rows=1, height=200)
+                        ds_subject = gr.Textbox(
+                            label="Subject description (used as caption prefix)",
+                            placeholder="e.g. a woman with red hair, green eyes, freckles",
+                            lines=2,
+                        )
+                        ds_extra = gr.Textbox(
+                            label="Extra prompt (appended to each pose)",
+                            placeholder="e.g. nude, studio lighting, white background",
+                            lines=1,
+                        )
+                        ds_count = gr.Slider(10, 150, value=50, step=5, label="Number of images")
+                        ds_lora = gr.Dropdown(LORA_TITLES, value="None (Base Klein 9B)", label="LoRA")
+                        ds_lora_str = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="LoRA Strength")
+                        with gr.Row():
+                            ds_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Starting Seed")
+                            ds_guidance = gr.Slider(0.0, 10.0, value=1.0, step=0.1, label="Guidance")
+                            ds_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                        ds_btn = gr.Button("Generate Dataset", variant="primary", size="lg")
+
+                    with gr.Column(scale=2):
+                        ds_gallery = gr.Gallery(label="Generated Dataset", columns=5, rows=3, height=500, object_fit="contain")
+                        ds_status = gr.Textbox(label="Status / Captions Preview", lines=8, interactive=False)
+
+                @spaces.GPU(duration=300)
+                def generate_dataset(ref_images, subject, extra, count, lora_name, lora_str,
+                                     seed, guidance, steps, progress=gr.Progress(track_tqdm=True)):
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                    try:
+                        pil_images = process_gallery_images(ref_images)
+                        if not pil_images:
+                            raise gr.Error("Upload a reference image!")
+
+                        ref = pil_images[0]
+                        count = int(count)
+                        poses = (POSE_LIBRARY * ((count // len(POSE_LIBRARY)) + 1))[:count]
+
+                        # LoRA
+                        style = LORA_MAP.get(lora_name)
+                        if style and style["adapter_name"]:
+                            try:
+                                pipe.load_lora_weights(style["repo"], weight_name=style["weights"],
+                                                       adapter_name=style["adapter_name"])
+                            except ValueError:
+                                pass
+                            pipe.set_adapters([style["adapter_name"]], adapter_weights=[lora_str])
+                        else:
+                            pipe.disable_lora()
+
+                        w, h = update_dimensions(ref)
+                        ref_resized = ref.resize((w, h), Image.LANCZOS).convert("RGB")
+
+                        results = []
+                        captions = []
+                        subject_text = subject.strip() if subject else "a person"
+                        extra_text = ", " + extra.strip() if extra and extra.strip() else ""
+
+                        for i, pose in enumerate(poses):
+                            progress((i + 1) / count, desc=f"Image {i+1}/{count}")
+                            caption = f"{subject_text}, {pose}{extra_text}"
+                            gen = torch.Generator(device=device).manual_seed(seed + i)
+                            img = pipe(
+                                image=ref_resized, prompt=caption,
+                                guidance_scale=guidance, width=w, height=h,
+                                num_inference_steps=steps, generator=gen,
+                            ).images[0]
+                            results.append((img, f"{i:03d}"))
+                            captions.append(f"{i:03d}.txt: {caption}")
+
+                        status = f"Generated {count} images.\n\nCaption preview:\n" + "\n".join(captions[:10])
+                        if count > 10:
+                            status += f"\n... and {count - 10} more"
+
+                        return results, status
+                    finally:
+                        gc.collect()
+                        torch.cuda.empty_cache()
+
+                ds_btn.click(
+                    fn=generate_dataset,
+                    inputs=[ds_ref, ds_subject, ds_extra, ds_count, ds_lora, ds_lora_str,
+                            ds_seed, ds_guidance, ds_steps],
+                    outputs=[ds_gallery, ds_status],
+                )
+
 if __name__ == "__main__":
     demo.queue().launch(ssr_mode=False, show_error=True)