rebuild: purpose-built tabs with optimal model/LoRA per task, Z-Image txt2img first tab

app.py

@@ -7,156 +7,109 @@ import torch
 import random
 from PIL import Image
 
+from diffusers import Flux2KleinPipeline, ZImagePipeline
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MAX_SEED = np.iinfo(np.int32).max
 
-# ---- Model Registry ----
-MODELS = {
-    "FLUX.2 Klein 9B (4 steps, editing)": {
-        "id": "black-forest-labs/FLUX.2-klein-9B",
-        "pipeline": "Flux2KleinPipeline",
-        "defaults": {"steps": 4, "guidance": 1.0},
-        "supports_strength": False,
-        "multi_image": True,
-    },
-    "Z-Image Turbo (9 steps, img2img)": {
-        "id": "Tongyi-MAI/Z-Image-Turbo",
-        "pipeline": "ZImageImg2ImgPipeline",
-        "defaults": {"steps": 9, "guidance": 0.0},
-        "supports_strength": True,
-        "multi_image": False,
-    },
-    "Qwen-Image-Edit 2511 (50 steps, editing)": {
-        "id": "Qwen/Qwen-Image-Edit-2511",
-        "pipeline": "QwenImageEditPlusPipeline",
-        "defaults": {"steps": 50, "guidance": 1.0},
-        "supports_strength": False,
-        "multi_image": True,
-    },
+# ---- 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"},
 }
-MODEL_NAMES = list(MODELS.keys())
-
-# ---- Prompt Templates ----
-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.
-FROM PICTURE 1 (strictly preserve):
-- Scene: lighting conditions, shadows, highlights, color temperature, environment, background
-- Head positioning: exact rotation angle, tilt, direction the head is facing
-- Expression: facial expression, micro-expressions, eye gaze direction, mouth position, emotion
-FROM PICTURE 2 (strictly preserve identity):
-- Facial structure: face shape, bone structure, jawline, chin
-- All facial features: eye color, eye shape, nose structure, lip shape and fullness, eyebrows
-- Hair: color, style, texture, hairline
-- Skin: texture, tone, complexion
-The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k."""
-
-PROMPT_TEMPLATES = {
-    "None": "",
-    "Face Swap (Detailed)": FACE_SWAP_PROMPT,
-    "Face Swap (Concise)": "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.",
-    "Expression Transfer": "h34d_sw4p: replace the head of Picture 1 by the head from Picture 2, strictly preserving the identity, facial features (eyes, nose, mouth), and skin texture of Picture 2. ensure the new head mimics the identical expression, angle, and rotation found in Picture 1.",
-    "Clothing Transfer": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions.",
-    "Style Transfer": "Transform this image into the artistic style shown, preserving the subject's identity and pose exactly. Match the color palette, brushwork, and aesthetic of the reference.",
-    "Pose Change - Sitting": "Change the person's pose to sitting comfortably in an armchair, preserving their identity and clothing exactly.",
-    "Pose Change - Side Profile": "Show the person from a side profile, looking to the right, preserving their identity and clothing exactly.",
-    "Pose Change - Arms Crossed": "Change the pose so the person has their arms crossed, looking confident. Preserve identity and clothing exactly.",
-    "Relight (Neutral Studio)": "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows. Maintain the original identity of all subjects exactly.",
-    "Remove Background": "Remove the background and replace it with a clean white studio backdrop. Keep the subject exactly as they are.",
-    "De-censor (Remove Bars)": "De-censor the image by removing black bars. Restore the original image content underneath naturally.",
-    "Upscale / Enhance": "Enhance this image to higher quality. Sharpen details, improve clarity, fix any compression artifacts. Preserve everything exactly as-is but at higher quality, 4k, sharp details.",
+
+# ---- 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.",
 }
 
-# Character sheet view prompts
-CHAR_SHEET_VIEWS = [
-    ("Front Face", "Show the person's face from directly in front, looking straight at the camera. Head and shoulders only. Preserve identity exactly. Clean background."),
-    ("Left Profile", "Show the person's face from the left side, 90 degree left profile view. Head and shoulders only. Preserve identity exactly. Clean background."),
-    ("Right Profile", "Show the person's face from the right side, 90 degree right profile view. Head and shoulders only. Preserve identity exactly. Clean background."),
-    ("Front Body", "Show the person's full body from directly in front, standing in a neutral pose. Preserve identity, clothing, and proportions exactly. Clean background."),
-    ("Left Body", "Show the person's full body from the left side, 90 degree profile. Standing in a neutral pose. Preserve identity, clothing, and proportions exactly. Clean background."),
-    ("Right Body", "Show the person's full body from the right side, 90 degree profile. Standing in a neutral pose. Preserve identity, clothing, and proportions exactly. Clean background."),
-    ("Back Body", "Show the person's full body from behind, back view. Standing in a neutral pose. Preserve identity, clothing, and proportions exactly. Clean background."),
+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",
 ]
 
-# ---- LoRA Registry ----
-LORA_STYLES = [
-    {
-        "title": "None (Base Klein 9B)",
-        "adapter_name": None,
-        "repo": None,
-        "weights": None,
-    },
-    {
-        "title": "BFS Head Swap V1 (rank64)",
-        "adapter_name": "bfs-head-v1-r64",
-        "repo": "Alissonerdx/BFS-Best-Face-Swap",
-        "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
-    },
-    {
-        "title": "BFS Head Swap V1 (rank128)",
-        "adapter_name": "bfs-head-v1-r128",
-        "repo": "Alissonerdx/BFS-Best-Face-Swap",
-        "weights": "bfs_head_v1_flux-klein_9b_step3500_rank128.safetensors",
-    },
-    {
-        "title": "Klein Consistency",
-        "adapter_name": "klein-consistency",
-        "repo": "dx8152/Flux2-Klein-9B-Consistency",
-        "weights": "Klein-consistency.safetensors",
-    },
-    {
-        "title": "Klein Delight Style",
-        "adapter_name": "klein-delight",
-        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
-        "weights": "pytorch_lora_weights.safetensors",
-    },
-    {
-        "title": "Ghost Mannequin",
-        "adapter_name": "ghost-mannequin",
-        "repo": "nhathoangfoto/FLUX.2-klein-ghost-mannequin",
-        "weights": "3D-GhosMannequinRank-256_000005000.safetensors",
-    },
-]
+# ---- 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!")
 
-LORA_TITLES = [s["title"] for s in LORA_STYLES]
-LORA_MAP = {s["title"]: s for s in LORA_STYLES}
-
-# ---- Model Manager ----
-_current_model = None
-_current_model_name = None
-pipe = None
-
-def load_model(model_name):
-    """Load a model by name. Swaps out the previous model to save VRAM."""
-    global pipe, _current_model_name
-    if model_name == _current_model_name and pipe is not None:
-        return pipe
-
-    # Unload previous
-    if pipe is not None:
-        del pipe
-        gc.collect()
-        torch.cuda.empty_cache()
-        print(f"Unloaded {_current_model_name}")
-
-    cfg = MODELS[model_name]
-    print(f"Loading {cfg['id']}...")
-
-    if cfg["pipeline"] == "Flux2KleinPipeline":
-        from diffusers import Flux2KleinPipeline
-        pipe = Flux2KleinPipeline.from_pretrained(cfg["id"], torch_dtype=torch.bfloat16).to(device)
-    elif cfg["pipeline"] == "ZImageImg2ImgPipeline":
-        from diffusers import ZImageImg2ImgPipeline
-        pipe = ZImageImg2ImgPipeline.from_pretrained(cfg["id"], torch_dtype=torch.bfloat16).to(device)
-    elif cfg["pipeline"] == "QwenImageEditPlusPipeline":
-        from diffusers import QwenImageEditPlusPipeline
-        pipe = QwenImageEditPlusPipeline.from_pretrained(cfg["id"], torch_dtype=torch.bfloat16).to(device)
-
-    _current_model_name = model_name
-    print(f"{model_name} loaded!")
-    return pipe
-
-# Pre-load Klein 9B at startup
-pipe = load_model("FLUX.2 Klein 9B (4 steps, editing)")
-LOADED_ADAPTERS = set()
+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!")
 
 
 # ---- Helpers ----
@@ -165,355 +118,361 @@ def update_dimensions(image):
         return 1024, 1024
     w, h = image.size
     scale = min(1024 / w, 1024 / h)
-    nw = (int(w * scale) // 16) * 16
-    nh = (int(h * scale) // 16) * 16
-    return nw, nh
+    return (int(w * scale) // 16) * 16, (int(h * scale) // 16) * 16
 
 
-def process_gallery_images(images):
+def process_images(images):
     if not images:
         return []
-    pil_images = []
+    out = []
     for item in images:
         try:
-            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
-            if isinstance(path_or_img, str):
-                pil_images.append(Image.open(path_or_img).convert("RGB"))
-            elif isinstance(path_or_img, Image.Image):
-                pil_images.append(path_or_img.convert("RGB"))
+            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:
-                pil_images.append(Image.open(path_or_img.name).convert("RGB"))
+                out.append(Image.open(p.name).convert("RGB"))
         except Exception as e:
-            print(f"Skipping image: {e}")
-    return pil_images
-
-
-def apply_template(template_name):
-    return PROMPT_TEMPLATES.get(template_name, "")
+            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 run_pipeline(model_name, images, prompt, guidance, steps, seed, strength=0.6):
-    """Run the active pipeline. Handles API differences between models."""
-    global pipe
-    pipe = load_model(model_name)
-    cfg = MODELS[model_name]
-    gen = torch.Generator(device=device).manual_seed(seed)
+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()
 
-    kwargs = dict(prompt=prompt, guidance_scale=guidance, width=w, height=h,
-                  num_inference_steps=steps, generator=gen)
-
-    if cfg["multi_image"]:
-        kwargs["image"] = processed if len(processed) > 1 else processed[0]
-    else:
-        kwargs["image"] = processed[0]
-
-    if cfg["supports_strength"]:
-        kwargs["strength"] = strength
-
-    return pipe(**kwargs).images[0]
-
-
-# ---- Inference ----
-def activate_lora(lora_name, lora_strength, custom_repo="", custom_file="", custom_strength=1.0):
-    """Load and activate LoRAs on the current pipeline."""
-    active_adapters = []
-    active_weights = []
-
-    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, AttributeError):
-            pass
-        active_adapters.append(style["adapter_name"])
-        active_weights.append(lora_strength)
 
-    if custom_repo and custom_repo.strip() and custom_file and custom_file.strip():
-        try:
-            pipe.load_lora_weights(custom_repo.strip(), weight_name=custom_file.strip(),
-                                   adapter_name="custom-lora")
-        except (ValueError, AttributeError):
-            pass
-        active_adapters.append("custom-lora")
-        active_weights.append(custom_strength)
+# ===========================================================
+# 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)
 
-    if active_adapters:
-        pipe.set_adapters(active_adapters, adapter_weights=active_weights)
-        print(f"Active LoRAs: {list(zip(active_adapters, active_weights))}")
-    else:
-        try:
-            pipe.disable_lora()
-        except AttributeError:
-            pass
+        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
+# ===========================================================
 @spaces.GPU
-def infer(
-    input_images, prompt, model_name, lora_name, lora_strength,
-    custom_lora_repo, custom_lora_file, custom_lora_strength,
-    seed, randomize_seed, guidance_scale, steps, strength,
-    progress=gr.Progress(track_tqdm=True),
-):
-    gc.collect()
-    torch.cuda.empty_cache()
-
+def image_edit(ref_images, prompt, nsfw_on, nsfw_str,
+               seed, randomize_seed, steps, guidance,
+               progress=gr.Progress(track_tqdm=True)):
+    gc.collect(); torch.cuda.empty_cache()
     try:
-        pil_images = process_gallery_images(input_images)
-        if not pil_images:
-            raise gr.Error("Upload at least one image!")
+        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 a prompt!")
+            raise gr.Error("Enter an edit prompt!")
 
-        global pipe
-        pipe = load_model(model_name)
-        activate_lora(lora_name, lora_strength, custom_lora_repo, custom_lora_file, custom_lora_strength)
+        loras = []
+        if nsfw_on:
+            loras.append(("nsfw", nsfw_str))
+        activate_loras(loras)
 
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
-        image = run_pipeline(model_name, pil_images, prompt, guidance_scale, steps, seed, strength)
-        return image, seed
+        result = generate(images, prompt.strip(), guidance, steps, seed)
+        return result, seed
+    finally:
+        gc.collect(); torch.cuda.empty_cache()
+
+
+# ===========================================================
+# Tab 3: Pose Variations
+# ===========================================================
+@spaces.GPU
+def pose_variations(ref_images, subject, extra, poses_selected, nsfw_on, nsfw_str,
+                    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 = [("consistency", 1.0)]
+        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 = []
+
+        for i, pose in enumerate(poses_selected):
+            progress((i + 1) / len(poses_selected), desc=f"Pose {i+1}/{len(poses_selected)}")
+            prompt = f"{subject_text}, {pose}{extra_text}"
+            img = generate(images, prompt, guidance, steps, seed + i)
+            results.append((img, pose[:50]))
 
+        return results
     finally:
-        gc.collect()
-        torch.cuda.empty_cache()
-
-
-POSE_LIBRARY = (
-    # 7 character sheet views
-    [v[1] for v in CHAR_SHEET_VIEWS] +
-    # 50 additional poses
-    [
-        "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",
-        "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",
-        "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",
-        "medium shot from waist up, one hand raised",
-        "full body shot, standing tall, power pose",
-    ]
-)
+        gc.collect(); torch.cuda.empty_cache()
 
 
+# ===========================================================
+# Tab 4: Dataset Generator
+# ===========================================================
 @spaces.GPU(duration=300)
-def generate_dataset(
-    ref_images, subject, extra, count, model_name, lora_name, lora_str,
-    seed, guidance, steps, strength,
-    progress=gr.Progress(track_tqdm=True),
-):
-    gc.collect()
-    torch.cuda.empty_cache()
+def generate_dataset(ref_images, subject, extra, count, nsfw_on, nsfw_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:
+        images = process_images(ref_images)
+        if not images:
             raise gr.Error("Upload at least one reference image!")
 
-        global pipe
-        pipe = load_model(model_name)
-        activate_lora(lora_name, lora_str)
+        loras = [("consistency", 1.0)]
+        if nsfw_on:
+            loras.append(("nsfw", nsfw_str))
+        activate_loras(loras)
 
         count = int(count)
         poses = (POSE_LIBRARY * ((count // len(POSE_LIBRARY)) + 1))[:count]
-
-        cfg = MODELS[model_name]
-        w, h = update_dimensions(pil_images[0])
-        refs = [img.resize((w, h), Image.LANCZOS).convert("RGB") for img in pil_images]
+        subject_text = subject.strip() if subject and subject.strip() else "a person"
+        extra_text = ", " + extra.strip() if extra and extra.strip() else ""
 
         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}"
-            image = run_pipeline(model_name, pil_images, caption, guidance, steps, seed + i, strength)
-            results.append((image, f"{i:03d}"))
+            img = generate(images, caption, guidance, steps, seed + i)
+            results.append((img, f"{i:03d}"))
             captions.append(f"{i:03d}.txt: {caption}")
 
-        status = f"Generated {count} images with {model_name}.\n\nCaption preview:\n" + "\n".join(captions[:10])
-        if count > 10:
-            status += f"\n... and {count - 10} more"
+        status = f"Generated {count} images.\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
     finally:
-        gc.collect()
-        torch.cuda.empty_cache()
+        gc.collect(); torch.cuda.empty_cache()
 
 
-# ---- UI ----
-css = "#col-container { margin: 0 auto; max-width: 1100px; }"
+# ===========================================================
+# UI
+# ===========================================================
+css = "#app { margin: 0 auto; max-width: 1100px; }"
 
 with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("# FLUX.2 Klein Studio\nMulti-model image editing: Klein 9B, Z-Image Turbo, Qwen-Image-Edit")
+    with gr.Column(elem_id="app"):
+        gr.Markdown("# FLUX.2 Klein Studio\nText prompt → Generate → Edit → Pose → LoRA Dataset. Full pipeline.")
 
         with gr.Tabs():
-            # ==================== IMAGE EDIT / FACE SWAP ====================
-            with gr.TabItem("Image Edit / Face Swap"):
+            # ==================== 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(scale=2):
-                        input_images = gr.Gallery(
-                            label="Input Images (1 for edit, 2 for face swap: body first, face second)",
-                            type="filepath", columns=2, rows=1, height=280, allow_preview=True,
-                        )
-                        prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the edit...")
+                    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():
-                            template_select = gr.Dropdown(list(PROMPT_TEMPLATES.keys()), value="None", label="Prompt Template", scale=2)
-                            run_btn = gr.Button("Generate", variant="primary", scale=1, size="lg")
-
-                        model_select = gr.Dropdown(MODEL_NAMES, value=MODEL_NAMES[0], label="Model")
-                        strength = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Strength (Z-Image Turbo only)",
-                                             info="How much to change from original. Ignored by Klein/Qwen.")
-
-                        with gr.Accordion("LoRA Adapters", open=True):
-                            lora_select = gr.Dropdown(LORA_TITLES, value="None (Base Klein 9B)", label="Built-in LoRA")
-                            lora_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="LoRA Strength")
-                            gr.Markdown("**Custom LoRA** (stacks with built-in)")
-                            with gr.Row():
-                                custom_repo = gr.Textbox(label="HF Repo", placeholder="user/my-lora", scale=2)
-                                custom_file = gr.Textbox(label="Filename", placeholder="weights.safetensors", scale=2)
-                                custom_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Strength", scale=1)
-
-                        with gr.Accordion("Generation Settings", open=False):
-                            seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
-                            randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
-                            with gr.Row():
-                                guidance = gr.Slider(0.0, 10.0, value=1.0, step=0.1, label="Guidance Scale")
-                                steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                            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")
 
-                    with gr.Column(scale=2):
-                        output_image = gr.Image(label="Result", interactive=False, format="png", height=500)
-                        seed_output = gr.Number(label="Seed Used", interactive=False)
-
-                template_select.change(fn=apply_template, inputs=[template_select], outputs=[prompt])
-                run_btn.click(
-                    fn=infer,
-                    inputs=[input_images, prompt, model_select, lora_select, lora_strength,
-                            custom_repo, custom_file, custom_strength,
-                            seed, randomize_seed, guidance, steps, strength],
-                    outputs=[output_image, seed_output],
-                )
-
-            # ==================== QUICK EDIT ====================
-            with gr.TabItem("Quick Image Edit"):
-                gr.Markdown("Upload one image + prompt. Simple and fast.")
+                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():
-                        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, add tattoos...")
-                        qe_model = gr.Dropdown(MODEL_NAMES, value=MODEL_NAMES[0], label="Model")
-                        qe_strength = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Strength (Z-Image only)")
-                        qe_lora = gr.Dropdown(LORA_TITLES, value="None (Base Klein 9B)", label="LoRA")
-                        qe_lora_str = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="LoRA Strength")
+                        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():
-                            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")
+                            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():
-                        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, model_name, img_strength, lora_name, lora_str,
-                               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!")
-                        global pipe
-                        pipe = load_model(model_name)
-                        activate_lora(lora_name, lora_str)
-                        if randomize:
-                            seed = random.randint(0, MAX_SEED)
-                        return run_pipeline(model_name, [image], prompt, guidance, steps, seed, img_strength), seed
-                    finally:
-                        gc.collect()
-                        torch.cuda.empty_cache()
-
-                qe_btn.click(
-                    fn=quick_edit,
-                    inputs=[qe_image, qe_prompt, qe_model, qe_strength, qe_lora, qe_lora_str,
-                            qe_steps, qe_guidance, qe_seed, qe_rand],
-                    outputs=[qe_output, qe_seed_out],
-                )
-
-            # ==================== DATASET GENERATOR (includes 360 sheet) ====================
-            with gr.TabItem("LoRA Dataset / 360 Sheet"):
-                gr.Markdown("Generate captioned pose datasets for LoRA training. First 7 images are 360 character sheet views, then 50 additional poses cycle. Multi-reference supported.")
+                        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("Upload image(s) and describe the edit. NSFW LoRA on by default.")
+                with gr.Row():
+                    with gr.Column():
+                        ie_images = gr.Gallery(label="Input Images", 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 prompt", lines=3, placeholder="Describe what to change...")
+                        with gr.Row():
+                            ie_nsfw = gr.Checkbox(value=True, label="NSFW LoRA")
+                            ie_nsfw_str = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="NSFW strength")
+                        with gr.Row():
+                            ie_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                            ie_guidance = gr.Slider(0.0, 10.0, value=1.0, step=0.1, label="Guidance")
+                        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_nsfw, ie_nsfw_str, ie_seed, ie_rand, ie_steps, ie_guidance],
+                    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)",
+                        )
+                        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_btn.click(fn=pose_variations,
+                    inputs=[pv_ref, pv_subject, pv_extra, pv_poses, pv_nsfw, pv_nsfw_str,
+                            pv_seed, pv_guidance, pv_steps],
+                    outputs=[pv_gallery])
+
+            # ==================== 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 description (caption prefix)", placeholder="e.g. a woman with red hair, green eyes", lines=2)
-                        ds_extra = gr.Textbox(label="Extra prompt (appended to each)", placeholder="e.g. nude, studio lighting, white background", lines=1)
-                        ds_count = gr.Slider(7, 150, value=50, step=1, label="Number of images (first 7 = 360 sheet)")
-                        ds_model = gr.Dropdown(MODEL_NAMES, value=MODEL_NAMES[0], label="Model")
-                        ds_strength = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Strength (Z-Image only)")
-                        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")
+                        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")
+                        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_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", lines=8, interactive=False)
-
-                ds_btn.click(
-                    fn=generate_dataset,
-                    inputs=[ds_ref, ds_subject, ds_extra, ds_count, ds_model, ds_lora, ds_lora_str,
-                            ds_seed, ds_guidance, ds_steps, ds_strength],
-                    outputs=[ds_gallery, ds_status],
-                )
+                        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_btn.click(fn=generate_dataset,
+                    inputs=[ds_ref, ds_subject, ds_extra, ds_count, ds_nsfw, ds_nsfw_str,
+                            ds_seed, ds_guidance, ds_steps],
+                    outputs=[ds_gallery, ds_status])
 
 if __name__ == "__main__":
     demo.queue().launch(ssr_mode=False, show_error=True)