import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
import tempfile
import zipfile
from io import BytesIO
from PIL import Image

from diffusers import Flux2KleinPipeline, ZImagePipeline, QwenImageEditPlusPipeline
from transformers import AutoProcessor, AutoModelForCausalLM

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_SEED = np.iinfo(np.int32).max

# ---- LoRA Config ----
LORAS = {
    "bfs-swap": {"repo": "Alissonerdx/BFS-Best-Face-Swap", "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors"},
    "nsfw": {"repo": "AntiLeecher/Flux-Klein-NSFW-Lora", "weights": "Flux Klein - NSFW v2.safetensors"},
    "consistency": {"repo": "dx8152/Flux2-Klein-9B-Consistency", "weights": "Klein-consistency.safetensors"},
    "delight": {"repo": "linoyts/Flux2-Klein-Delight-LoRA", "weights": "pytorch_lora_weights.safetensors"},
}

# ---- Prompt Presets ----
FACE_SWAP_PROMPT = (
    "head_swap: start with Picture 1 as the base image, keeping its lighting, "
    "environment, and background. Remove the head from Picture 1 completely and "
    "replace it with the head from Picture 2, strictly preserving the hair, eye color, "
    "nose structure of Picture 2. copy the direction of the eye, head rotation, "
    "micro expressions from Picture 1, high quality, sharp details, 4k."
)

EDIT_TEMPLATES = {
    "Custom": "",
    "Remove clothing": "Remove all clothing from the person. Artistic nudity, full body visible, photorealistic, sharp details.",
    "Change outfit": "Change the person's outfit to: ",
    "Add tattoos": "Add detailed tattoos covering the person's arms and torso. Preserve identity and pose exactly.",
    "Change hair": "Change the person's hairstyle to: ",
    "Remove background": "Remove the background and replace with a clean white studio backdrop.",
    "Relight (studio)": "Relight with neutral, uniform studio illumination. Soft, evenly distributed lighting. Preserve identity exactly.",
    "Age up": "Make the person appear 20 years older. Preserve identity, add wrinkles, grey hair, aged skin naturally.",
    "Age down": "Make the person appear 15 years younger. Preserve identity, smoother skin, more youthful features.",
    "De-censor": "De-censor the image by removing black bars and mosaic censoring. Restore the original image content underneath naturally.",
    "Enhance / Upscale": "Enhance this image to higher quality. Sharpen details, improve clarity, 4k, sharp details.",
}

POSE_LIBRARY = [
    # 7 character sheet views first
    "face from directly in front, looking straight at camera, head and shoulders, clean background",
    "face from left side, 90 degree left profile, head and shoulders, clean background",
    "face from right side, 90 degree right profile, head and shoulders, clean background",
    "full body from directly in front, standing neutral pose, clean background",
    "full body from left side, 90 degree profile, standing neutral, clean background",
    "full body from right side, 90 degree profile, standing neutral, clean background",
    "full body from behind, back view, standing neutral, clean background",
    # Additional poses
    "standing facing camera, neutral pose, arms at sides",
    "standing with arms crossed, confident pose",
    "standing with hands on hips",
    "standing three-quarter view from the left",
    "standing three-quarter view from the right",
    "standing side profile, looking right",
    "standing from behind, back view",
    "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",
    "kneeling on one knee",
    "kneeling on both knees, upright",
    "leaning against a wall, arms crossed",
    "leaning against a wall, one foot up",
    "walking towards camera, mid-stride",
    "walking away from camera, back view",
    "looking up at the sky, chin raised",
    "looking down, contemplative",
    "head tilted to the left, slight smile",
    "laughing naturally, candid expression",
    "hands behind head, stretching",
    "one hand touching hair, casual",
    "hands in pockets, casual standing",
    "arms raised above head, celebratory",
    "crouching down, low angle",
    "bending forward, looking at camera",
    "twisting torso, looking over shoulder",
    "dancing pose, one leg lifted",
    "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",
    "medium shot from waist up, arms at sides",
    "full body shot, standing tall, power pose",
    "sitting sideways on chair, arm draped over backrest",
    "leaning forward with hands on knees",
    "running towards camera, dynamic pose",
    "head tilted to the right, serious expression",
    "waving at camera, friendly gesture",
]

# ---- Load Models ----
print("Loading FLUX.2 Klein 9B...")
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B", torch_dtype=torch.bfloat16,
).to(device)
print("Klein loaded!")

print("Loading Z-Image Turbo...")
zimage_pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16,
).to(device)
print("Z-Image Turbo loaded!")

print("Loading Qwen-Image-Edit 2511...")
qwen_pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511", torch_dtype=torch.bfloat16,
).to(device)
print("Qwen loaded!")

print("Loading Florence-2 captioner...")
florence_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True, revision="refs/pr/6")
florence_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Florence-2-base-ft", torch_dtype=torch.float32, trust_remote_code=True, revision="refs/pr/6",
).to("cpu")
print("Florence-2 loaded (CPU)!")


# ---- Helpers ----
def update_dimensions(image):
    if image is None:
        return 1024, 1024
    w, h = image.size
    scale = min(1024 / w, 1024 / h)
    return (int(w * scale) // 16) * 16, (int(h * scale) // 16) * 16


def process_images(images):
    if not images:
        return []
    out = []
    for item in images:
        try:
            p = item[0] if isinstance(item, (tuple, list)) else item
            if isinstance(p, str):
                out.append(Image.open(p).convert("RGB"))
            elif isinstance(p, Image.Image):
                out.append(p.convert("RGB"))
            else:
                out.append(Image.open(p.name).convert("RGB"))
        except Exception as e:
            print(f"Skip: {e}")
    return out


def activate_loras(names_and_weights):
    """Activate a set of LoRAs by name. names_and_weights = [(name, weight), ...]"""
    active = []
    weights = []
    for name, w in names_and_weights:
        if name not in LORAS:
            continue
        cfg = LORAS[name]
        try:
            pipe.load_lora_weights(cfg["repo"], weight_name=cfg["weights"], adapter_name=name)
        except ValueError:
            pass  # already loaded
        active.append(name)
        weights.append(w)

    if active:
        pipe.set_adapters(active, adapter_weights=weights)
        print(f"LoRAs: {list(zip(active, weights))}")
    else:
        try:
            pipe.disable_lora()
        except Exception:
            pass


def caption_image(pil_image, prefix=""):
    """Generate a detailed caption for an image using Florence-2."""
    task = "<MORE_DETAILED_CAPTION>"
    inputs = florence_processor(text=task, images=pil_image, return_tensors="pt").to("cpu")
    with torch.no_grad():
        generated = florence_model.generate(
            **inputs, max_new_tokens=256, num_beams=3, early_stopping=True,
        )
    raw = florence_processor.batch_decode(generated, skip_special_tokens=False)[0]
    caption = florence_processor.post_process_generation(raw, task=task, image_size=pil_image.size)
    text = caption.get(task, "").strip()
    if prefix:
        text = f"{prefix}, {text}"
    return text


def generate(images, prompt, guidance, steps, seed):
    w, h = update_dimensions(images[0])
    processed = [img.resize((w, h), Image.LANCZOS).convert("RGB") for img in images]
    image_input = processed if len(processed) > 1 else processed[0]
    return pipe(
        image=image_input, prompt=prompt,
        guidance_scale=guidance, width=w, height=h,
        num_inference_steps=steps,
        generator=torch.Generator(device=device).manual_seed(seed),
    ).images[0]


# ===========================================================
# Tab 0: Text to Image (Z-Image Turbo)
# ===========================================================
@spaces.GPU
def txt2img(prompt, negative_prompt, seed, randomize_seed, steps, guidance, width, height,
            progress=gr.Progress(track_tqdm=True)):
    gc.collect(); torch.cuda.empty_cache()
    try:
        if not prompt or not prompt.strip():
            raise gr.Error("Enter a prompt!")
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
        result = zimage_pipe(
            prompt=prompt.strip(),
            negative_prompt=negative_prompt.strip() if negative_prompt else None,
            width=width, height=height,
            num_inference_steps=steps,
            guidance_scale=guidance,
            generator=torch.Generator(device=device).manual_seed(seed),
        ).images[0]
        return result, seed
    finally:
        gc.collect(); torch.cuda.empty_cache()


# ===========================================================
# Tab 1: Face Swap
# ===========================================================
@spaces.GPU
def face_swap(body_img, face_img, custom_prompt, nsfw_on, nsfw_str, swap_str,
              seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    gc.collect(); torch.cuda.empty_cache()
    try:
        body_images = process_images(body_img)
        face_images = process_images(face_img)
        if not body_images:
            raise gr.Error("Upload a body/scene image!")
        if not face_images:
            raise gr.Error("Upload a face reference image!")

        loras = [("bfs-swap", swap_str)]
        if nsfw_on:
            loras.append(("nsfw", nsfw_str))
        activate_loras(loras)

        prompt = custom_prompt.strip() if custom_prompt.strip() else FACE_SWAP_PROMPT
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        images = body_images + face_images
        result = generate(images, prompt, 1.0, 4, seed)
        return result, seed
    finally:
        gc.collect(); torch.cuda.empty_cache()


# ===========================================================
# Tab 2: Image Edit (Qwen-Image-Edit 2511)
# ===========================================================
@spaces.GPU(duration=120)
def image_edit(ref_images, prompt, seed, randomize_seed, steps,
               progress=gr.Progress(track_tqdm=True)):
    gc.collect(); torch.cuda.empty_cache()
    try:
        images = process_images(ref_images)
        if not images:
            raise gr.Error("Upload an image!")
        if not prompt or not prompt.strip():
            raise gr.Error("Enter an edit prompt!")

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        w, h = update_dimensions(images[0])
        processed = [img.resize((w, h), Image.LANCZOS).convert("RGB") for img in images]

        result = qwen_pipe(
            image=processed,
            prompt=prompt.strip(),
            num_inference_steps=steps,
            generator=torch.Generator(device=device).manual_seed(seed),
        ).images[0]
        return result, seed
    finally:
        gc.collect(); torch.cuda.empty_cache()


# ===========================================================
# Tab 3: Pose Variations
# ===========================================================
@spaces.GPU(duration=180)
def pose_variations(ref_images, subject, extra, poses_selected, nsfw_on, nsfw_str,
                    auto_caption, seed, guidance, steps,
                    progress=gr.Progress(track_tqdm=True)):
    gc.collect(); torch.cuda.empty_cache()
    try:
        images = process_images(ref_images)
        if not images:
            raise gr.Error("Upload a reference image!")
        if not poses_selected:
            raise gr.Error("Select at least one pose!")

        loras = []
        if nsfw_on:
            loras.append(("nsfw", nsfw_str))
        activate_loras(loras)

        subject_text = subject.strip() if subject and subject.strip() else "the person"
        extra_text = ", " + extra.strip() if extra and extra.strip() else ""
        results = []
        captions = []
        pil_results = []
        total = len(poses_selected)

        for i, pose in enumerate(poses_selected):
            progress((i + 1) / total, desc=f"Generating {i+1}/{total}")
            prompt = f"{subject_text}, {pose}{extra_text}"
            img = generate(images, prompt, guidance, steps, seed + i)

            if auto_caption:
                progress((i + 1) / total, desc=f"Captioning {i+1}/{total}")
                caption = caption_image(img, prefix=subject_text)
            else:
                caption = prompt

            results.append((img, pose[:50]))
            pil_results.append((img, caption))
            captions.append(f"{i:03d}: {caption}")

        # Build ZIP
        zip_path = tempfile.mktemp(suffix=".zip")
        with zipfile.ZipFile(zip_path, "w") as zf:
            for i, (img, caption) in enumerate(pil_results):
                buf = BytesIO()
                img.save(buf, format="PNG")
                zf.writestr(f"{i:03d}.png", buf.getvalue())
                zf.writestr(f"{i:03d}.txt", caption)

        cap_type = "Florence-2" if auto_caption else "prompt-based"
        status = f"{total} poses, {cap_type} captions.\n\n" + "\n".join(captions[:10])
        if total > 10:
            status += f"\n... +{total - 10} more"

        return results, status, zip_path
    finally:
        gc.collect(); torch.cuda.empty_cache()


# ===========================================================
# Tab 4: Dataset Generator
# ===========================================================
@spaces.GPU(duration=600)
def generate_dataset(ref_images, subject, extra, count, nsfw_on, nsfw_str,
                     auto_caption, seed, guidance, steps,
                     progress=gr.Progress(track_tqdm=True)):
    gc.collect(); torch.cuda.empty_cache()
    try:
        images = process_images(ref_images)
        if not images:
            raise gr.Error("Upload at least one reference image!")

        loras = []
        if nsfw_on:
            loras.append(("nsfw", nsfw_str))
        activate_loras(loras)

        count = int(count)
        poses = (POSE_LIBRARY * ((count // len(POSE_LIBRARY)) + 1))[:count]
        subject_text = subject.strip() if subject and subject.strip() else "a person"
        extra_text = ", " + extra.strip() if extra and extra.strip() else ""

        results = []
        captions = []
        pil_results = []

        for i, pose in enumerate(poses):
            progress((i + 1) / count, desc=f"Generating {i+1}/{count}")
            gen_prompt = f"{subject_text}, {pose}{extra_text}"
            img = generate(images, gen_prompt, guidance, steps, seed + i)

            # Caption: use Florence-2 or fall back to generation prompt
            if auto_caption:
                progress((i + 1) / count, desc=f"Captioning {i+1}/{count}")
                caption = caption_image(img, prefix=subject_text)
            else:
                caption = gen_prompt

            results.append((img, f"{i:03d}"))
            pil_results.append((img, caption))
            captions.append(f"{i:03d}.txt: {caption}")

        # Build ZIP with images + caption .txt files
        zip_path = tempfile.mktemp(suffix=".zip")
        with zipfile.ZipFile(zip_path, "w") as zf:
            for i, (img, caption) in enumerate(pil_results):
                buf = BytesIO()
                img.save(buf, format="PNG")
                zf.writestr(f"{i:03d}.png", buf.getvalue())
                zf.writestr(f"{i:03d}.txt", caption)

        cap_type = "Florence-2 auto-caption" if auto_caption else "prompt-based caption"
        status = f"Generated {count} images with {cap_type}.\nFirst 7 = 360 character sheet views.\n\n"
        status += "Caption preview:\n" + "\n".join(captions[:15])
        if count > 15:
            status += f"\n... +{count - 15} more"
        return results, status, zip_path
    finally:
        gc.collect(); torch.cuda.empty_cache()


# ===========================================================
# UI
# ===========================================================
css = "#app { margin: 0 auto; max-width: 1100px; }"

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="app"):
        gr.Markdown("# FLUX.2 Klein Studio\nText prompt → Generate → Edit → Pose → LoRA Dataset. Full pipeline.")

        with gr.Tabs():
            # ==================== TEXT TO IMAGE ====================
            with gr.TabItem("Text to Image"):
                gr.Markdown("Generate from a text prompt using Z-Image Turbo. No censorship. Use the output as a starting point for the other tabs.")
                with gr.Row():
                    with gr.Column():
                        t2i_prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the character/scene...")
                        t2i_neg = gr.Textbox(label="Negative prompt", lines=1, value="worst quality, low quality, blurry, deformed")
                        with gr.Row():
                            t2i_w = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                            t2i_h = gr.Slider(512, 1536, value=1024, step=64, label="Height")
                        with gr.Row():
                            t2i_steps = gr.Slider(1, 20, value=9, step=1, label="Steps")
                            t2i_guidance = gr.Slider(0.0, 10.0, value=0.0, step=0.1, label="Guidance (0 for Turbo)")
                        t2i_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                        t2i_rand = gr.Checkbox(value=True, label="Randomize seed")
                        t2i_btn = gr.Button("Generate", variant="primary", size="lg")
                    with gr.Column():
                        t2i_out = gr.Image(label="Result", interactive=False, format="png", height=500)
                        t2i_seed_out = gr.Number(label="Seed")

                t2i_btn.click(fn=txt2img,
                    inputs=[t2i_prompt, t2i_neg, t2i_seed, t2i_rand, t2i_steps, t2i_guidance, t2i_w, t2i_h],
                    outputs=[t2i_out, t2i_seed_out])

            # ==================== FACE SWAP ====================
            with gr.TabItem("Face Swap"):
                gr.Markdown("Upload body/scene as Picture 1, face reference as Picture 2. BFS Head Swap LoRA auto-loaded.")
                with gr.Row():
                    with gr.Column():
                        fs_body = gr.Gallery(label="Body / Scene (Picture 1)", type="filepath", columns=1, rows=1, height=220)
                        fs_face = gr.Gallery(label="Face Reference (Picture 2)", type="filepath", columns=1, rows=1, height=220)
                        fs_prompt = gr.Textbox(label="Custom prompt (leave empty for default swap prompt)", lines=2)
                        with gr.Row():
                            fs_nsfw = gr.Checkbox(value=True, label="NSFW LoRA")
                            fs_nsfw_str = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="NSFW strength")
                            fs_swap_str = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Swap strength")
                        fs_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                        fs_rand = gr.Checkbox(value=True, label="Randomize seed")
                        fs_btn = gr.Button("Swap Faces", variant="primary", size="lg")
                    with gr.Column():
                        fs_out = gr.Image(label="Result", interactive=False, format="png", height=500)
                        fs_seed_out = gr.Number(label="Seed")

                fs_btn.click(fn=face_swap,
                    inputs=[fs_body, fs_face, fs_prompt, fs_nsfw, fs_nsfw_str, fs_swap_str, fs_seed, fs_rand],
                    outputs=[fs_out, fs_seed_out])

            # ==================== IMAGE EDIT ====================
            with gr.TabItem("Image Edit"):
                gr.Markdown("Powered by **Qwen-Image-Edit 2511** — instruction-based editing at 50 steps. Supports multi-reference (upload multiple images). No LoRA needed.")
                with gr.Row():
                    with gr.Column():
                        ie_images = gr.Gallery(label="Input Images (multi-reference supported)", type="filepath", columns=2, rows=1, height=280)
                        ie_template = gr.Dropdown(list(EDIT_TEMPLATES.keys()), value="Custom", label="Preset")
                        ie_prompt = gr.Textbox(label="Edit instruction", lines=3, placeholder="e.g. remove clothing, add tattoos, change background...")
                        ie_steps = gr.Slider(10, 100, value=50, step=5, label="Steps (50 recommended)")
                        ie_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                        ie_rand = gr.Checkbox(value=True, label="Randomize seed")
                        ie_btn = gr.Button("Edit", variant="primary", size="lg")
                    with gr.Column():
                        ie_out = gr.Image(label="Result", interactive=False, format="png", height=500)
                        ie_seed_out = gr.Number(label="Seed")

                ie_template.change(fn=lambda t: EDIT_TEMPLATES.get(t, ""), inputs=[ie_template], outputs=[ie_prompt])
                ie_btn.click(fn=image_edit,
                    inputs=[ie_images, ie_prompt, ie_seed, ie_rand, ie_steps],
                    outputs=[ie_out, ie_seed_out])

            # ==================== POSE VARIATIONS ====================
            with gr.TabItem("Pose Variations"):
                gr.Markdown("Generate the same character in different poses. Consistency + NSFW LoRAs auto-loaded.")
                with gr.Row():
                    with gr.Column(scale=1):
                        pv_ref = gr.Gallery(label="Reference Images", type="filepath", columns=2, rows=1, height=200)
                        pv_subject = gr.Textbox(label="Subject description", placeholder="e.g. a woman with red hair", lines=1)
                        pv_extra = gr.Textbox(label="Extra prompt (appended to each)", placeholder="e.g. nude, studio lighting", lines=1)
                        pv_poses = gr.CheckboxGroup(
                            choices=POSE_LIBRARY[:20],  # Show first 20 for selection
                            value=POSE_LIBRARY[:7],  # Default: 360 sheet views
                            label="Select poses (first 7 = 360 character sheet)",
                        )
                        pv_autocap = gr.Checkbox(value=True, label="Auto-caption with Florence-2")
                        with gr.Row():
                            pv_nsfw = gr.Checkbox(value=True, label="NSFW LoRA")
                            pv_nsfw_str = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="NSFW strength")
                        with gr.Row():
                            pv_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                            pv_guidance = gr.Slider(0.0, 10.0, value=1.0, step=0.1, label="Guidance")
                            pv_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
                        pv_btn = gr.Button("Generate Poses", variant="primary", size="lg")
                    with gr.Column(scale=2):
                        pv_gallery = gr.Gallery(label="Results", columns=4, rows=2, height=500, object_fit="contain")
                        pv_status = gr.Textbox(label="Captions", lines=6, interactive=False)
                        pv_zip = gr.File(label="Download ZIP (images + captions)")

                pv_btn.click(fn=pose_variations,
                    inputs=[pv_ref, pv_subject, pv_extra, pv_poses, pv_nsfw, pv_nsfw_str,
                            pv_autocap, pv_seed, pv_guidance, pv_steps],
                    outputs=[pv_gallery, pv_status, pv_zip])

            # ==================== DATASET GENERATOR ====================
            with gr.TabItem("LoRA Dataset"):
                gr.Markdown("Batch-generate captioned images for LoRA training. First 7 = 360 sheet, then cycles through 50 poses. Consistency + NSFW LoRAs auto-loaded.")
                with gr.Row():
                    with gr.Column(scale=1):
                        ds_ref = gr.Gallery(label="Reference Images", type="filepath", columns=2, rows=1, height=200)
                        ds_subject = gr.Textbox(label="Subject (caption prefix)", placeholder="e.g. a woman with red hair, green eyes, freckles", lines=2)
                        ds_extra = gr.Textbox(label="Extra (appended to each caption)", placeholder="e.g. nude, studio lighting, white background", lines=1)
                        ds_count = gr.Slider(7, 150, value=50, step=1, label="Number of images")
                        ds_autocap = gr.Checkbox(value=True, label="Auto-caption with Florence-2",
                                                 info="Describes what's actually in each image. Better for LoRA training than prompt-based captions.")
                        with gr.Row():
                            ds_nsfw = gr.Checkbox(value=True, label="NSFW LoRA")
                            ds_nsfw_str = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="NSFW 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="Dataset", columns=5, rows=3, height=500, object_fit="contain")
                        ds_status = gr.Textbox(label="Captions", lines=8, interactive=False)
                        ds_zip = gr.File(label="Download ZIP (images + captions)")

                ds_btn.click(fn=generate_dataset,
                    inputs=[ds_ref, ds_subject, ds_extra, ds_count, ds_nsfw, ds_nsfw_str,
                            ds_autocap, ds_seed, ds_guidance, ds_steps],
                    outputs=[ds_gallery, ds_status, ds_zip])

if __name__ == "__main__":
    demo.queue().launch(ssr_mode=False, show_error=True)