Update app.py

app.py

@@ -1,660 +1,1176 @@
 import os
+import gc
+import torch
+import gradio as gr
+from PIL import Image
+from pathlib import Path
+import shutil
 import json
-import copy
-import time
-import requests
-import random
-import logging
+from huggingface_hub import hf_hub_download, HfApi, create_repo, upload_file
+import tempfile
+import multiprocessing as mp
+
+# Training imports
+from peft import LoraConfig, get_peft_model
+from tqdm.auto import tqdm
 import numpy as np
-import spaces
-from typing import Any, Dict, List, Optional, Union
 
-import torch
-from PIL import Image
-import gradio as gr
+# Set memory allocation config BEFORE any CUDA operations
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,garbage_collection_threshold:0.6'
+
+# Global state
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+# HF Token from environment
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Model repo
+MODEL_REPO = "Tongyi-MAI/Z-Image-Turbo"
+
+
+# ============================================
+# Comic Style CSS
+# ============================================
+COMIC_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');
+
+.gradio-container {
+    background-color: #FEF9C3 !important;
+    background-image: radial-gradient(#1F2937 1px, transparent 1px) !important;
+    background-size: 20px 20px !important;
+    min-height: 100vh !important;
+    font-family: 'Comic Neue', cursive, sans-serif !important;
+}
+
+footer, .footer, .gradio-container footer, .built-with, [class*="footer"], .gradio-footer, a[href*="gradio.app"] {
+    display: none !important;
+    visibility: hidden !important;
+    height: 0 !important;
+}
+
+/* HOME Button Style */
+.home-button-container {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    gap: 15px;
+    margin-bottom: 15px;
+    padding: 12px 20px;
+    background: linear-gradient(135deg, #10B981 0%, #059669 100%);
+    border: 4px solid #1F2937;
+    border-radius: 12px;
+    box-shadow: 6px 6px 0 #1F2937;
+}
+
+.home-button {
+    display: inline-flex;
+    align-items: center;
+    gap: 8px;
+    padding: 10px 25px;
+    background: linear-gradient(135deg, #FACC15 0%, #F59E0B 100%);
+    color: #1F2937;
+    font-family: 'Bangers', cursive;
+    font-size: 1.4rem;
+    letter-spacing: 2px;
+    text-decoration: none;
+    border: 3px solid #1F2937;
+    border-radius: 8px;
+    box-shadow: 4px 4px 0 #1F2937;
+    transition: all 0.2s ease;
+}
+
+.home-button:hover {
+    background: linear-gradient(135deg, #FDE047 0%, #FACC15 100%);
+    transform: translate(-2px, -2px);
+    box-shadow: 6px 6px 0 #1F2937;
+}
+
+.home-button:active {
+    transform: translate(2px, 2px);
+    box-shadow: 2px 2px 0 #1F2937;
+}
+
+.url-display {
+    font-family: 'Comic Neue', cursive;
+    font-size: 1.1rem;
+    font-weight: 700;
+    color: #FFF;
+    background: rgba(0,0,0,0.3);
+    padding: 8px 16px;
+    border-radius: 6px;
+    border: 2px solid rgba(255,255,255,0.3);
+}
+
+.header-container {
+    text-align: center;
+    padding: 25px 20px;
+    background: linear-gradient(135deg, #3B82F6 0%, #8B5CF6 100%);
+    border: 4px solid #1F2937;
+    border-radius: 12px;
+    margin-bottom: 20px;
+    box-shadow: 8px 8px 0 #1F2937;
+    position: relative;
+}
+
+.header-title {
+    font-family: 'Bangers', cursive !important;
+    color: #FFF !important;
+    font-size: 2.8rem !important;
+    text-shadow: 3px 3px 0 #1F2937 !important;
+    letter-spacing: 3px !important;
+    margin: 0 !important;
+}
+
+.header-subtitle {
+    font-family: 'Comic Neue', cursive !important;
+    font-size: 1.1rem !important;
+    color: #FEF9C3 !important;
+    margin-top: 8px !important;
+    font-weight: 700 !important;
+}
+
+.stats-badge {
+    display: inline-block;
+    background: #FACC15;
+    color: #1F2937;
+    padding: 6px 14px;
+    border-radius: 20px;
+    font-size: 0.9rem;
+    margin: 3px;
+    font-weight: 700;
+    border: 2px solid #1F2937;
+    box-shadow: 2px 2px 0 #1F2937;
+}
+
+.gr-panel, .gr-box, .gr-form, .block, .gr-group {
+    background: #FFF !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    box-shadow: 5px 5px 0 #1F2937 !important;
+}
+
+.gr-button-primary, button.primary, .gr-button.primary {
+    background: linear-gradient(135deg, #EF4444 0%, #F97316 100%) !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    color: #FFF !important;
+    font-family: 'Bangers', cursive !important;
+    font-size: 1.3rem !important;
+    letter-spacing: 2px !important;
+    padding: 12px 24px !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+    text-shadow: 1px 1px 0 #1F2937 !important;
+    transition: all 0.2s ease !important;
+}
+
+.gr-button-primary:hover, button.primary:hover {
+    background: linear-gradient(135deg, #DC2626 0%, #EA580C 100%) !important;
+    transform: translate(-2px, -2px) !important;
+    box-shadow: 6px 6px 0 #1F2937 !important;
+}
+
+.gr-button-primary:active, button.primary:active {
+    transform: translate(2px, 2px) !important;
+    box-shadow: 2px 2px 0 #1F2937 !important;
+}
+
+textarea, input[type="text"], input[type="number"] {
+    background: #FFF !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    color: #1F2937 !important;
+    font-family: 'Comic Neue', cursive !important;
+    font-weight: 700 !important;
+}
+
+textarea:focus, input[type="text"]:focus {
+    border-color: #3B82F6 !important;
+    box-shadow: 3px 3px 0 #3B82F6 !important;
+}
+
+.info-box {
+    background: linear-gradient(135deg, #FACC15 0%, #FDE047 100%) !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    padding: 12px 15px !important;
+    margin: 10px 0 !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+    font-family: 'Comic Neue', cursive !important;
+    font-weight: 700 !important;
+    color: #1F2937 !important;
+}
+
+.result-box textarea {
+    background: #1F2937 !important;
+    color: #10B981 !important;
+    font-family: 'Courier New', monospace !important;
+    border: 3px solid #10B981 !important;
+    border-radius: 8px !important;
+    box-shadow: 4px 4px 0 #10B981 !important;
+}
+
+label, .gr-input-label, .gr-block-label {
+    color: #1F2937 !important;
+    font-family: 'Comic Neue', cursive !important;
+    font-weight: 700 !important;
+}
+
+.gr-accordion {
+    background: #E0F2FE !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+}
+
+.tab-nav button {
+    font-family: 'Comic Neue', cursive !important;
+    font-weight: 700 !important;
+    border: 2px solid #1F2937 !important;
+    margin: 2px !important;
+}
+
+.tab-nav button.selected {
+    background: #3B82F6 !important;
+    color: #FFF !important;
+    box-shadow: 3px 3px 0 #1F2937 !important;
+}
+
+.footer-comic {
+    text-align: center;
+    padding: 20px;
+    background: linear-gradient(135deg, #3B82F6 0%, #8B5CF6 100%);
+    border: 4px solid #1F2937;
+    border-radius: 12px;
+    margin-top: 20px;
+    box-shadow: 6px 6px 0 #1F2937;
+}
+
+.footer-comic p {
+    font-family: 'Comic Neue', cursive !important;
+    color: #FFF !important;
+    margin: 5px 0 !important;
+    font-weight: 700 !important;
+}
+
+::-webkit-scrollbar {
+    width: 12px;
+    height: 12px;
+}
+
+::-webkit-scrollbar-track {
+    background: #FEF9C3;
+    border: 2px solid #1F2937;
+}
+
+::-webkit-scrollbar-thumb {
+    background: #3B82F6;
+    border: 2px solid #1F2937;
+    border-radius: 6px;
+}
+
+::-webkit-scrollbar-thumb:hover {
+    background: #EF4444;
+}
+
+::selection {
+    background: #FACC15;
+    color: #1F2937;
+}
+
+/* Slider Styling */
+input[type="range"] {
+    accent-color: #3B82F6;
+}
+
+.gr-slider input[type="range"]::-webkit-slider-thumb {
+    background: #EF4444 !important;
+    border: 2px solid #1F2937 !important;
+}
+
+/* Image/Gallery Container */
+.gr-image, .gr-gallery {
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+}
+
+/* Quality Badge */
+.quality-badge {
+    display: inline-block;
+    background: linear-gradient(135deg, #10B981 0%, #059669 100%);
+    color: white;
+    padding: 4px 12px;
+    border-radius: 15px;
+    font-size: 0.8rem;
+    font-weight: bold;
+    border: 2px solid #1F2937;
+    margin-left: 8px;
+}
+
+#col-container {
+    max-width: 1200px;
+    margin: 0 auto;
+}
+
+/* Hide Hugging Face elements */
+.huggingface-space-link,
+a[href*="huggingface.co/spaces"],
+button[class*="share"],
+.share-button,
+[class*="hf-logo"],
+.gr-share-btn,
+#hf-logo,
+.hf-icon,
+svg[class*="hf"],
+div[class*="huggingface"],
+a[class*="huggingface"],
+.svelte-1rjryqp,
+header a[href*="huggingface"],
+.space-header,
+div.absolute.right-0 a[href*="huggingface"],
+.gr-group > a[href*="huggingface"],
+a[target="_blank"][href*="huggingface.co"] {
+    display: none !important;
+    visibility: hidden !important;
+    opacity: 0 !important;
+    pointer-events: none !important;
+    width: 0 !important;
+    height: 0 !important;
+    overflow: hidden !important;
+}
+
+/* Training specific styles */
+.training-section {
+    background: linear-gradient(135deg, #E0F2FE 0%, #DBEAFE 100%) !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 12px !important;
+    padding: 15px !important;
+    margin: 10px 0 !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+}
+
+.tips-box {
+    background: linear-gradient(135deg, #D1FAE5 0%, #A7F3D0 100%) !important;
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    padding: 12px 15px !important;
+    margin: 10px 0 !important;
+    box-shadow: 4px 4px 0 #1F2937 !important;
+    font-family: 'Comic Neue', cursive !important;
+    font-weight: 700 !important;
+    color: #1F2937 !important;
+}
+"""
 
-from diffusers import (
-    DiffusionPipeline,
-    AutoencoderKL,
-    ZImagePipeline
-)
-
-from huggingface_hub import (
-    hf_hub_download,
-    HfFileSystem,
-    ModelCard,
-    snapshot_download)
-
-from diffusers.utils import load_image
-from typing import Iterable
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
-
-colors.orange_red = colors.Color(
-    name="orange_red",
-    c50="#FFF0E5",
-    c100="#FFE0CC",
-    c200="#FFC299",
-    c300="#FFA366",
-    c400="#FF8533",
-    c500="#FF4500",
-    c600="#E63E00",
-    c700="#CC3700",
-    c800="#B33000",
-    c900="#992900",
-    c950="#802200",
-)
-
-class OrangeRedTheme(Soft):
-    def __init__(
-        self,
-        *,
-        primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red, # Use the new color
-        neutral_hue: colors.Color | str = colors.slate,
-        text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
-        ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            background_fill_primary="*primary_50",
-            background_fill_primary_dark="*primary_900",
-            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
-            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="white",
-            button_primary_text_color_hover="white",
-            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_secondary_text_color="black",
-            button_secondary_text_color_hover="white",
-            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
-            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
-            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
-            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
-            slider_color="*secondary_500",
-            slider_color_dark="*secondary_600",
-            block_title_text_weight="600",
-            block_border_width="3px",
-            block_shadow="*shadow_drop_lg",
-            button_primary_shadow="*shadow_drop_lg",
-            button_large_padding="11px",
-            color_accent_soft="*primary_100",
-            block_label_background_fill="*primary_200",
-        )
 
-orange_red_theme = OrangeRedTheme()
+def aggressive_cleanup():
+    """Aggressively clean up GPU memory"""
+    gc.collect()
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.reset_accumulated_memory_stats()
 
-loras = [
-    # 로컬 jimin LoRA (app.py와 같은 디렉토리에 jimin.safetensors 필요)
-    {
-        "image": "https://i.namu.wiki/i/umL8EZtn0hs-nMRYeFxIrkGrMe-R1u5c9fJE8ufrLjvXz52VcSIbG7TT9QJoL2rR7vsFww1lLrE4bwfn5uOBzfq9a90HGdNdlTLmr_KoqOchTovbVC3RDzhDbp7FI-Wq-esCu7_BYIptqethL4onBg.webp",
+
+def get_gpu_memory_info():
+    """Get current GPU memory status"""
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated(0) / 1e9
+        reserved = torch.cuda.memory_reserved(0) / 1e9
+        total = torch.cuda.get_device_properties(0).total_memory / 1e9
+        free = total - allocated
+        return f"GPU: {allocated:.1f}GB allocated, {reserved:.1f}GB reserved, {free:.1f}GB free of {total:.1f}GB"
+    return "No GPU"
+
+
+def check_gpu():
+    """Check GPU availability and memory"""
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
+        return f"✅ GPU: {gpu_name} ({gpu_mem:.1f}GB total)"
+    return "❌ No GPU detected"
+
+
+def check_hf_token():
+    """Check if HF_TOKEN is configured"""
+    if HF_TOKEN:
+        try:
+            api = HfApi(token=HF_TOKEN)
+            user_info = api.whoami()
+            return f"✅ Logged in as: {user_info['name']}"
+        except Exception as e:
+            return f"⚠️ Token invalid: {str(e)}"
+    return "❌ HF_TOKEN not set"
+
+
+def get_hf_username():
+    if HF_TOKEN:
+        try:
+            api = HfApi(token=HF_TOKEN)
+            return api.whoami()['name']
+        except:
+            return None
+    return None
+
+
+# ============================================
+# SUBPROCESS-BASED CAPTIONING
+# ============================================
+
+def _caption_worker(image_paths_queue, results_queue, trigger_word, is_person):
+    """Worker process for Florence-2 captioning - completely isolated GPU context"""
+    import torch
+    from PIL import Image
+    
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    try:
+        from transformers import AutoProcessor, AutoModelForCausalLM
+        
+        print("[Caption Worker] Loading Florence-2-large...")
+        processor = AutoProcessor.from_pretrained(
+            "microsoft/Florence-2-large",
+            trust_remote_code=True
+        )
+        # FIX: Add attn_implementation="eager" to avoid _supports_sdpa error
+        model = AutoModelForCausalLM.from_pretrained(
+            "microsoft/Florence-2-large",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            attn_implementation="eager"  # Disable SDPA to avoid attribute error
+        ).to(device)
+        model.eval()
+        print("[Caption Worker] Florence-2 loaded!")
         
-        "title": "jimin Style",
-        "repo": "./",  # 로컬 경로
-        "weights": "jimin.safetensors",
-        "trigger_word": "jimin"    
-    },
-    {
-        "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/1111111111.png",
-        "title": "AWPortrait Z",
-        "repo": "Shakker-Labs/AWPortrait-Z", #1
-        "weights": "AWPortrait-Z.safetensors",
-        "trigger_word": "Portrait"    
-    },
-
-    {
-        "image": "https://cdn-uploads.huggingface.co/production/uploads/653cd3049107029eb004f968/DLCGlF9uUnFo5zxR5uyx6.png",
-        "title": "50s Western",
-        "repo": "neph1/50s_western_lora_zit", 
-        "weights": "50s_western_z_100.safetensors",
-        "trigger_word": "50s_western"    
-    },
-    {
-        "image": "https://huggingface.co/Sumitc13/Z-image-Turbo_LogC4_lora/resolve/main/images/1764464517272__000005000_1.jpg",
-        "title": "LogC4",
-        "repo": "Sumitc13/Z-image-Turbo_LogC4_lora", #30
-        "weights": "z-image-logc4_000005000.safetensors",
-        "trigger_word": "LogC4"    
-    },
-    {
-        "image": "https://huggingface.co/neph1/80s_scifi_lora_zit/resolve/main/images/ComfyUI_10288_.png",
-        "title": "80s Scifi",
-        "repo": "neph1/80s_scifi_lora_zit", 
-        "weights": "80s_scifi_z_80.safetensors",
-        "trigger_word": "80s_scifi"    
-    },
-
- # --------------------------------------------------------------------------------------------------------------------------------------   
-    {
-        "image": "https://huggingface.co/Ttio2/Z-Image-Turbo-pencil-sketch/resolve/main/images/z-image_00097_.png",
-        "title": "Turbo Pencil",
-        "repo": "Ttio2/Z-Image-Turbo-pencil-sketch", #0
-        "weights": "Zimage_pencil_sketch.safetensors",
-        "trigger_word": "pencil sketch"    
-    },
-    {
-        "image": "https://huggingface.co/neph1/50s_scifi_lora_zit/resolve/main/images/ComfyUI_08067_.png",
-        "title": "50s Scifi",
-        "repo": "neph1/50s_scifi_lora_zit",
-        "weights": "50s_scifi_z_80.safetensors",
-        "trigger_word": "50s_scifi"    
-    },
-    {
-        "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/cookie-mons.png",
-        "title": "Yarn Art Style",
-        "repo": "linoyts/yarn-art-style", #28
-        "weights": "yarn-art-style_000001250.safetensors",
-        "trigger_word": "yarn art style"    
-    },
-    {
-        "image": "https://huggingface.co/Quorlen/Z-Image-Turbo-Behind-Reeded-Glass-Lora/resolve/main/images/ComfyUI_00391_.png",
-        "title": "Behind Reeded Glass",
-        "repo": "Quorlen/Z-Image-Turbo-Behind-Reeded-Glass-Lora", #26
-        "weights": "Z_Image_Turbo_Behind_Reeded_Glass_Lora_TAV2_000002750.safetensors",
-        "trigger_word": "Act1vate!, Behind reeded glass"    
-    },
-    {
-        "image": "https://huggingface.co/ostris/z_image_turbo_childrens_drawings/resolve/main/images/1764433619736__000003000_9.jpg",
-        "title": "Childrens Drawings",
-        "repo": "ostris/z_image_turbo_childrens_drawings", #2
-        "weights": "z_image_turbo_childrens_drawings.safetensors",
-        "trigger_word": "Children Drawings"    
-    },
-    {
-        "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/xcxc.png",
-        "title": "Tarot Z",
-        "repo": "multimodalart/tarot-z-image-lora", #22
-        "weights": "tarot-z-image_000001250.safetensors",
-        "trigger_word": "trtcrd"    
-    },
-    {
-        "image": "https://huggingface.co/renderartist/Technically-Color-Z-Image-Turbo/resolve/main/images/ComfyUI_00917_.png",
-        "title": "Technically Color Z",
-        "repo": "renderartist/Technically-Color-Z-Image-Turbo", #3
-        "weights": "Technically_Color_Z_Image_Turbo_v1_renderartist_2000.safetensors",
-        "trigger_word": "t3chnic4lly"    
-    },
-    {
-        "image": "https://huggingface.co/SkyAsl/Tattoo-artist-Z/resolve/main/images/a%20dragon%20with%20flames.png",
-        "title": "Tattoo-artist-Z",
-        "repo": "SkyAsl/Tattoo-artist-Z", #31
-        "weights": "adapter_model.safetensors",
-        "trigger_word": "a tattoo design"    
-    },
-    {
-        "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/z-image_00147_.png",
-        "title": "Turbo Ghibli",
-        "repo": "Ttio2/Z-Image-Turbo-Ghibli-Style", #19
-        "weights": "ghibli_zimage_finetune.safetensors",
-        "trigger_word": "Ghibli Style"    
-    },
-    {
-        "image": "https://huggingface.co/tarn59/pixel_art_style_lora_z_image_turbo/resolve/main/images/ComfyUI_00273_.png",
-        "title": "Pixel Art",
-        "repo": "tarn59/pixel_art_style_lora_z_image_turbo", #4
-        "weights": "pixel_art_style_z_image_turbo.safetensors",
-        "trigger_word": "Pixel art style."    
-    },
-    {
-        "image": "https://huggingface.co/renderartist/Saturday-Morning-Z-Image-Turbo/resolve/main/images/Saturday_Morning_Z_15.png",
-        "title": "Saturday Morning",
-        "repo": "renderartist/Saturday-Morning-Z-Image-Turbo", #5
-        "weights": "Saturday_Morning_Z_Image_Turbo_v1_renderartist_1250.safetensors",
-        "trigger_word": "saturd4ym0rning"    
-    },
-    {
-        "image": "https://huggingface.co/AIImageStudio/ReversalFilmGravure_z_Image_turbo/resolve/main/images/2025-12-01_173047-z_image_z_image_turbo_bf16-435125750859057-euler_10_hires.png",
-        "title": "ReversalFilmGravure",
-        "repo": "AIImageStudio/ReversalFilmGravure_z_Image_turbo", #6
-        "weights": "z_image_turbo_ReversalFilmGravure_v1.0.safetensors",
-        "trigger_word": "Reversal Film Gravure, analog film photography"    
-    },
-    {
-        "image": "https://huggingface.co/renderartist/Coloring-Book-Z-Image-Turbo-LoRA/resolve/main/images/CBZ_00274_.png",
-        "title": "Coloring Book Z",
-        "repo": "renderartist/Coloring-Book-Z-Image-Turbo-LoRA", #7
-        "weights": "Coloring_Book_Z_Image_Turbo_v1_renderartist_2000.safetensors",
-        "trigger_word": "c0l0ringb00k"    
-    },
-    {
-        "image": "https://huggingface.co/damnthatai/1950s_American_Dream/resolve/main/images/ZImage_20251129163459_135x_00001_.jpg",
-        "title": "1950s American Dream",
-        "repo": "damnthatai/1950s_American_Dream", #8
-        "weights": "5os4m3r1c4n4_z.safetensors",
-        "trigger_word": "5os4m3r1c4n4, 1950s, painting, a painting of"    
-    },
-    {
-        "image": "https://huggingface.co/wcde/Z-Image-Turbo-DeJPEG-Lora/resolve/main/images/01.png",
-        "title": "DeJPEG",
-        "repo": "wcde/Z-Image-Turbo-DeJPEG-Lora", #9
-        "weights": "dejpeg_v3.safetensors",
-        "trigger_word": ""    
-    },
-    {
-        "image": "https://huggingface.co/suayptalha/Z-Image-Turbo-Realism-LoRA/resolve/main/images/n4aSpqa-YFXYo4dtcIg4W.png",
-        "title": "DeJPEG",
-        "repo": "suayptalha/Z-Image-Turbo-Realism-LoRA", #10
-        "weights": "pytorch_lora_weights.safetensors",
-        "trigger_word": "Realism"    
-    },
-    {
-        "image": "https://huggingface.co/renderartist/Classic-Painting-Z-Image-Turbo-LoRA/resolve/main/images/Classic_Painting_Z_00247_.png",
-        "title": "Classic Painting Z",
-        "repo": "renderartist/Classic-Painting-Z-Image-Turbo-LoRA", #11
-        "weights": "Classic_Painting_Z_Image_Turbo_v1_renderartist_1750.safetensors",
-        "trigger_word": "class1cpa1nt"    
-    },
-    {
-        "image": "https://huggingface.co/DK9/3D_MMORPG_style_z-image-turbo_lora/resolve/main/images/10_with_lora.png",
-        "title": "3D MMORPG",
-        "repo": "DK9/3D_MMORPG_style_z-image-turbo_lora", #12
-        "weights": "lostark_v1.safetensors",
-        "trigger_word": ""    
-    },
-    {
-        "image": "https://huggingface.co/Danrisi/Olympus_UltraReal_ZImage/resolve/main/images/Z-Image_01011_.png",
-        "title": "Olympus UltraReal",
-        "repo": "Danrisi/Olympus_UltraReal_ZImage", #13
-        "weights": "Olympus.safetensors",
-        "trigger_word": "digital photography, early 2000s compact camera aesthetic, amateur candid shot, digital photography, early 2000s compact camera aesthetic, amateur candid shot, direct flash lighting, hard flash shadow, specular highlights, overexposed highlights"    
-    },
-    {
-        "image": "https://huggingface.co/AiAF/D-ART_Z-Image-Turbo_LoRA/resolve/main/images/example_l3otpwzaz.png",
-        "title": "D ART Z Image",
-        "repo": "AiAF/D-ART_Z-Image-Turbo_LoRA", #14
-        "weights": "D-ART_Z-Image-Turbo.safetensors",
-        "trigger_word": "D-ART"    
-    },
-    {
-        "image": "https://huggingface.co/AlekseyCalvin/Marionette_Modernism_Z-image-Turbo_LoRA/resolve/main/bluebirdmandoll.webp",
-        "title": "Marionette Modernism",
-        "repo": "AlekseyCalvin/Marionette_Modernism_Z-image-Turbo_LoRA", #15
-        "weights": "ZImageDadadoll_000003600.safetensors",
-        "trigger_word": "DADADOLL style"    
-    },
-    {
-        "image": "https://huggingface.co/AlekseyCalvin/HistoricColor_Z-image-Turbo-LoRA/resolve/main/HSTZgen2.webp",
-        "title": "Historic Color Z",
-        "repo": "AlekseyCalvin/HistoricColor_Z-image-Turbo-LoRA", #16
-        "weights": "ZImage1HST_000004000.safetensors",
-        "trigger_word": "HST style"    
-    },
-    {
-        "image": "https://huggingface.co/tarn59/80s_air_brush_style_z_image_turbo/resolve/main/images/ComfyUI_00707_.png",
-        "title": "80s Air Brush",
-        "repo": "tarn59/80s_air_brush_style_z_image_turbo", #17
-        "weights": "80s_air_brush_style_v2_z_image_turbo.safetensors",
-        "trigger_word": "80s Air Brush style."    
-    },
-    {
-        "image": "https://huggingface.co/CedarC/Z-Image_360/resolve/main/images/1765505225357__000006750_6.jpg",
-        "title": "360panorama",
-        "repo": "CedarC/Z-Image_360", #18
-        "weights": "Z-Image_360.safetensors",
-        "trigger_word": "360panorama"    
-    },
-    {
-        "image": "https://huggingface.co/HAV0X1014/Z-Image-Turbo-KF-Bat-Eared-Fox-LoRA/resolve/main/images/ComfyUI_00132_.png",
-        "title": "KF-Bat-Eared",
-        "repo": "HAV0X1014/Z-Image-Turbo-KF-Bat-Eared-Fox-LoRA", #21
-        "weights": "z-image-turbo-bat_eared_fox.safetensors",
-        "trigger_word": "bat_eared_fox_kemono_friends"    
-    },
-    {
-        "image": "https://cdn-uploads.huggingface.co/production/uploads/653cd3049107029eb004f968/IHttgddXu6ZBMo7eyy8p6.png",
-        "title": "80s Horror",
-        "repo": "neph1/80s_horror_movies_lora_zit", #23
-        "weights": "80s_horror_z_80.safetensors",
-        "trigger_word": "80s_horror"    
-    },
-    {
-        "image": "https://huggingface.co/Quorlen/z_image_turbo_Sunbleached_Protograph_Style_Lora/resolve/main/images/ComfyUI_00024_.png",
-        "title": "Sunbleached Protograph",
-        "repo": "Quorlen/z_image_turbo_Sunbleached_Protograph_Style_Lora", #24
-        "weights": "zimageturbo_Sunbleach_Photograph_Style_Lora_TAV2_000002750.safetensors",
-        "trigger_word": "Act1vate!"    
-    },
-    {
-        "image": "https://huggingface.co/bunnycore/Z-Art-2.1/resolve/main/images/ComfyUI_00069_.png",
-        "title": "Z-Art-2.1",
-        "repo": "bunnycore/Z-Art-2.1", #25
-        "weights": "Z-Image-Art2.1.safetensors",
-        "trigger_word": "anime art"    
-    },
-    {
-        "image": "https://huggingface.co/cactusfriend/longfurby-z/resolve/main/images/1764658860954__000003000_1.jpg",
-        "title": "Longfurby",
-        "repo": "cactusfriend/longfurby-z", #27
-        "weights": "longfurbyZ.safetensors",
-        "trigger_word": ""    
-    },
-    {
-        "image": "https://huggingface.co/SkyAsl/Pixel-artist-Z/resolve/main/pixel-art-result.png",
-        "title": "Pixel Art",
-        "repo": "SkyAsl/Pixel-artist-Z", #29
-        "weights": "adapter_model.safetensors",
-        "trigger_word": "a pixel art character"    
-    },
-]
-
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "Tongyi-MAI/Z-Image-Turbo"
-
-print(f"Loading {base_model} pipeline...")
-
-# Initialize Pipeline
-pipe = ZImagePipeline.from_pretrained(
-    base_model,
-    torch_dtype=dtype,
-    low_cpu_mem_usage=False,
-).to(device)
-
-# ======== AoTI compilation + FA3 ========
-# As per reference for optimization
-try:
-    print("Applying AoTI compilation and FA3...")
-    pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-    spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
-    print("Optimization applied successfully.")
-except Exception as e:
-    print(f"Optimization warning: {e}. Continuing with standard pipeline.")
-
-MAX_SEED = np.iinfo(np.int32).max
-
-class calculateDuration:
-    def __init__(self, activity_name=""):
-        self.activity_name = activity_name
-
-    def __enter__(self):
-        self.start_time = time.time()
-        return self
+        image_paths = image_paths_queue.get()
+        captions = []
+        
+        for idx, img_path in enumerate(image_paths):
+            try:
+                img = Image.open(img_path).convert("RGB")
+                task = "<DETAILED_CAPTION>"
+                inputs = processor(text=task, images=img, return_tensors="pt")
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+                
+                with torch.no_grad():
+                    generated_ids = model.generate(
+                        **inputs, max_new_tokens=256, num_beams=3, do_sample=False
+                    )
+                
+                generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+                parsed = processor.post_process_generation(
+                    generated_text, task=task, image_size=(img.width, img.height)
+                )
+                
+                raw_caption = parsed.get(task, "")
+                
+                if raw_caption:
+                    caption = raw_caption.strip()
+                    if is_person:
+                        replacements = [
+                            "a young woman", "a woman", "the woman", "a young man", "a man", "the man",
+                            "a person", "the person", "a young person", "an individual",
+                            "a girl", "the girl", "a boy", "the boy",
+                            "a lady", "the lady", "a gentleman", "the gentleman",
+                            "someone", "a figure", "the figure"
+                        ]
+                        caption_lower = caption.lower()
+                        replaced = False
+                        for ref in replacements:
+                            if ref in caption_lower:
+                                import re
+                                pattern = re.compile(re.escape(ref), re.IGNORECASE)
+                                caption = pattern.sub(trigger_word, caption, count=1)
+                                replaced = True
+                                break
+                        if not replaced:
+                            caption = f"{trigger_word}, {caption}"
+                    else:
+                        caption = f"{trigger_word}, {caption}"
+                else:
+                    caption = trigger_word
+                
+                captions.append(caption)
+                print(f"[Caption Worker] [{idx+1}/{len(image_paths)}] {caption[:80]}...")
+                
+                del inputs, generated_ids, img
+                torch.cuda.empty_cache()
+                
+            except Exception as e:
+                print(f"[Caption Worker] Error on image {idx}: {e}")
+                captions.append(trigger_word)
+        
+        results_queue.put(captions)
+        del model, processor
+        torch.cuda.empty_cache()
+        
+    except Exception as e:
+        print(f"[Caption Worker] Fatal error: {e}")
+        import traceback
+        traceback.print_exc()
+        # Return trigger word as fallback - need to get image_paths first
+        try:
+            image_paths = image_paths_queue.get(timeout=1)
+            results_queue.put([trigger_word] * len(image_paths))
+        except:
+            results_queue.put([trigger_word] * 20)
+
+
+def run_captioning_subprocess(image_paths, trigger_word, is_person=True):
+    """Run captioning in subprocess for complete GPU memory isolation"""
+    print(f"[Main] Starting captioning subprocess...")
+    print(f"[Main] Before: {get_gpu_memory_info()}")
     
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.end_time = time.time()
-        self.elapsed_time = self.end_time - self.start_time
-        if self.activity_name:
-            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
-        else:
-            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
-
-def update_selection(evt: gr.SelectData, width, height):
-    selected_lora = loras[evt.index]
-    new_placeholder = f"Type a prompt for {selected_lora['title']}"
-    lora_repo = selected_lora["repo"]
-    # 로컬 LoRA 처리
-    if lora_repo == "./":
-        updated_text = f"### Selected: Local LoRA - {selected_lora['title']} ✅"
-    else:
-        updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
-    if "aspect" in selected_lora:
-        if selected_lora["aspect"] == "portrait":
-            width = 768
-            height = 1024
-        elif selected_lora["aspect"] == "landscape":
-            width = 1024
-            height = 768
-        else:
-            width = 1024
-            height = 1024
-    return (
-        gr.update(placeholder=new_placeholder),
-        updated_text,
-        evt.index,
-        width,
-        height,
+    ctx = mp.get_context('spawn')
+    image_paths_queue = ctx.Queue()
+    results_queue = ctx.Queue()
+    
+    worker = ctx.Process(
+        target=_caption_worker,
+        args=(image_paths_queue, results_queue, trigger_word, is_person)
     )
+    worker.start()
+    image_paths_queue.put(image_paths)
+    
+    try:
+        captions = results_queue.get(timeout=900)  # 15 min timeout
+    except Exception as e:
+        print(f"[Main] Captioning error: {e}")
+        captions = [trigger_word] * len(image_paths)
+    
+    worker.join(timeout=30)
+    if worker.is_alive():
+        worker.terminate()
+        worker.join()
+    
+    print(f"[Main] After captioning: {get_gpu_memory_info()}")
+    return captions
+
 
-@spaces.GPU
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
-    # Clean up previous LoRAs in both cases
-    with calculateDuration("Unloading LoRA"):
-        pipe.unload_lora_weights()
+def prepare_dataset(images, trigger_word, output_dir, use_auto_caption=True, is_person=True):
+    """Prepare dataset with subprocess captioning"""
+    dataset_dir = Path(output_dir) / "dataset"
+    dataset_dir.mkdir(parents=True, exist_ok=True)
+
+    image_paths = []
     
-    # Check if a LoRA is selected
-    if selected_index is not None and selected_index < len(loras):
-        selected_lora = loras[selected_index]
-        lora_path = selected_lora["repo"]
-        trigger_word = selected_lora["trigger_word"]
-        
-        # Prepare Prompt with Trigger Word
-        if(trigger_word):
-            if "trigger_position" in selected_lora:
-                if selected_lora["trigger_position"] == "prepend":
-                    prompt_mash = f"{trigger_word} {prompt}"
-                else:
-                    prompt_mash = f"{prompt} {trigger_word}"
-            else:
-                prompt_mash = f"{trigger_word} {prompt}"
+    for i, img in enumerate(images):
+        if img is None:
+            continue
+        if isinstance(img, tuple):
+            img = img[0]
+        if isinstance(img, str):
+            img_pil = Image.open(img)
+        elif isinstance(img, np.ndarray):
+            img_pil = Image.fromarray(img)
+        elif hasattr(img, 'mode'):
+            img_pil = img
         else:
-            prompt_mash = prompt
+            continue
 
-        # Load LoRA
-        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-            weight_name = selected_lora.get("weights", None)
-            try:
-                pipe.load_lora_weights(
-                    lora_path, 
-                    weight_name=weight_name, 
-                    adapter_name="default",
-                    low_cpu_mem_usage=True
-                )
-                # Set adapter scale
-                pipe.set_adapters(["default"], adapter_weights=[lora_scale])
-            except Exception as e:
-                print(f"Error loading LoRA: {e}")
-                gr.Warning("Failed to load LoRA weights. Generating with base model.")
+        if img_pil.mode != "RGB":
+            img_pil = img_pil.convert("RGB")
+
+        img_path = dataset_dir / f"image_{i:04d}.jpg"
+        img_pil.save(img_path, quality=95)
+        image_paths.append(str(img_path))
+    
+    if use_auto_caption:
+        captions = run_captioning_subprocess(image_paths, trigger_word, is_person)
     else:
-        # Base Model Case
-        print("No LoRA selected. Running with Base Model.")
-        prompt_mash = prompt
-        
-    with calculateDuration("Randomizing seed"):
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
+        captions = [trigger_word] * len(image_paths)
     
-    generator = torch.Generator(device=device).manual_seed(seed)
+    for img_path, caption in zip(image_paths, captions):
+        caption_path = Path(img_path).with_suffix('.txt')
+        caption_path.write_text(caption)
+    
+    return image_paths, captions, str(dataset_dir)
+
+
+def compute_flow_matching_loss(model_output, target, timesteps):
+    """Compute Rectified Flow loss"""
+    loss = torch.nn.functional.mse_loss(model_output, target, reduction="none")
+    loss = loss.mean(dim=list(range(1, len(loss.shape))))
+    return loss.mean()
+
 
-    # Note: Z-Image-Turbo is strictly T2I in this reference implementation. 
-    # Img2Img via image_input is disabled/ignored for this pipeline update.
+def upload_to_hub(lora_path, repo_name, trigger_word, training_info, progress_callback=None):
+    """Upload to HF Hub"""
+    if not HF_TOKEN:
+        return False, "HF_TOKEN not configured"
     
-    with calculateDuration("Generating image"):
-        # For Turbo models, guidance_scale is typically 0.0
-        forced_guidance = 0.0 # Turbo mode
+    try:
+        api = HfApi(token=HF_TOKEN)
+        username = get_hf_username()
+        if not username:
+            return False, "Could not get username"
         
-        final_image = pipe(
-            prompt=prompt_mash,
-            height=int(height),
-            width=int(width),
-            num_inference_steps=int(steps),
-            guidance_scale=forced_guidance,
-            generator=generator,
-        ).images[0]
+        repo_id = f"{username}/{repo_name}"
         
-    yield final_image, seed, gr.update(visible=False)
-
-def get_huggingface_safetensors(link):
-  split_link = link.split("/")
-  if(len(split_link) == 2):
-            model_card = ModelCard.load(link)
-            base_model = model_card.data.get("base_model")
-            print(base_model)
-      
-            # Relaxed check to allow Z-Image or Flux or others, assuming user knows what they are doing
-            # or specifically check for Z-Image-Turbo
-            if base_model not in ["Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.1-dev"]:
-                # Just a warning instead of error to allow experimentation
-                print("Warning: Base model might not match.")
-                
-            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-            trigger_word = model_card.data.get("instance_prompt", "")
-            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-            fs = HfFileSystem()
-            try:
-                list_of_files = fs.ls(link, detail=False)
-                for file in list_of_files:
-                    if(file.endswith(".safetensors")):
-                        safetensors_name = file.split("/")[-1]
-                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                      image_elements = file.split("/")
-                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-            except Exception as e:
-              print(e)
-              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            return split_link[1], link, safetensors_name, trigger_word, image_url
-
-def check_custom_model(link):
-    if(link.startswith("https://")):
-        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
-            link_split = link.split("huggingface.co/")
-            return get_huggingface_safetensors(link_split[1])
-    else: 
-        return get_huggingface_safetensors(link)
-
-def add_custom_lora(custom_lora):
-    global loras
-    if(custom_lora):
         try:
-            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
-            print(f"Loaded custom LoRA: {repo}")
-            card = f'''
-            <div class="custom_lora_card">
-              <span>Loaded custom LoRA:</span>
-              <div class="card_internal">
-                <img src="{image}" />
-                <div>
-                    <h3>{title}</h3>
-                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
-                </div>
-              </div>
-            </div>
-            '''
-            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
-            if(not existing_item_index):
-                new_item = {
-                    "image": image,
-                    "title": title,
-                    "repo": repo,
-                    "weights": path,
-                    "trigger_word": trigger_word
-                }
-                print(new_item)
-                existing_item_index = len(loras)
-                loras.append(new_item)
+            create_repo(repo_id=repo_id, token=HF_TOKEN, private=True, repo_type="model", exist_ok=True)
+        except:
+            pass
         
-            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
+        api.upload_file(
+            path_or_fileobj=lora_path,
+            path_in_repo=f"{repo_name}.safetensors",
+            repo_id=repo_id,
+            token=HF_TOKEN
+        )
+        
+        readme = f"""---
+license: apache-2.0
+base_model: Tongyi-MAI/Z-Image-Turbo
+tags: [lora, z-image, text-to-image, diffusers]
+---
+# {repo_name}
+Trigger: `{trigger_word}`
+{training_info}
+"""
+        api.upload_file(path_or_fileobj=readme.encode(), path_in_repo="README.md", repo_id=repo_id, token=HF_TOKEN)
+        
+        return True, f"https://huggingface.co/{repo_id}"
+    except Exception as e:
+        return False, str(e)
+
+
+def train_lora(
+    images, trigger_word, output_name, num_steps, learning_rate, lora_rank,
+    resolution, batch_size, upload_to_hub_flag, hub_repo_name,
+    use_auto_caption, is_person_training, progress=gr.Progress()
+):
+    """Train LoRA with CORRECT ZImageTransformer2DModel forward signature"""
+
+    if not torch.cuda.is_available():
+        return None, "❌ No GPU available"
+
+    if not images or len(images) < 3:
+        return None, "❌ Please upload at least 3 images"
+
+    if not trigger_word:
+        return None, "❌ Please specify a trigger word"
+
+    if not output_name:
+        output_name = "z_image_lora"
+    output_name = output_name.replace(" ", "_").lower()
+
+    if upload_to_hub_flag and not HF_TOKEN:
+        return None, "❌ HF_TOKEN not configured"
+
+    progress(0, desc="Initializing...")
+    print(f"[Train] Start: {get_gpu_memory_info()}")
+    aggressive_cleanup()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        try:
+            # ============================================
+            # PHASE 1: Captioning (Subprocess)
+            # ============================================
+            progress(0.02, desc="Running Florence-2 captioning (subprocess)...")
+            
+            image_paths, captions, dataset_dir = prepare_dataset(
+                images, trigger_word, tmpdir, use_auto_caption, is_person_training
+            )
+
+            if len(image_paths) < 3:
+                return None, "❌ Not enough valid images"
+
+            progress(0.12, desc=f"Captioning done: {len(image_paths)} images")
+            aggressive_cleanup()
+            print(f"[Train] After captioning cleanup: {get_gpu_memory_info()}")
+
+            # ============================================
+            # PHASE 2: Load Pipeline for Text Encoding
+            # ============================================
+            progress(0.15, desc="Loading pipeline for encoding...")
+            print(f"[Train] Before pipeline: {get_gpu_memory_info()}")
+            
+            from diffusers import ZImagePipeline
+            
+            # Load pipeline to CPU first
+            pipe = ZImagePipeline.from_pretrained(
+                MODEL_REPO,
+                torch_dtype=DTYPE,
+            )
+            
+            # Get VAE scaling factor
+            vae_scaling_factor = pipe.vae.config.scaling_factor
+            
+            # ============================================
+            # PHASE 3: Encode Captions with Text Encoder
+            # ============================================
+            progress(0.20, desc="Encoding captions...")
+            
+            # Move text encoder to GPU
+            pipe.text_encoder.to(DEVICE)
+            
+            cached_text_embeddings = []
+            
+            with torch.no_grad():
+                for idx, caption in enumerate(captions):
+                    text_inputs = pipe.tokenizer(
+                        caption,
+                        padding="max_length",
+                        max_length=256,
+                        truncation=True,
+                        return_tensors="pt"
+                    ).to(DEVICE)
+                    
+                    text_emb = pipe.text_encoder(**text_inputs)[0]
+                    cached_text_embeddings.append(text_emb.cpu())
+                    
+                    del text_inputs, text_emb
+                    
+                    if idx % 2 == 0:
+                        torch.cuda.empty_cache()
+                        progress(0.20 + 0.10 * (idx / len(captions)), 
+                                desc=f"Encoding caption {idx+1}/{len(captions)}")
+            
+            # Free text encoder
+            pipe.text_encoder.to("cpu")
+            del pipe.text_encoder
+            aggressive_cleanup()
+            print(f"[Train] After text encoding: {get_gpu_memory_info()}")
+
+            # ============================================
+            # PHASE 4: Encode Images with VAE
+            # ============================================
+            progress(0.32, desc="Encoding images with VAE...")
+            
+            pipe.vae.to(DEVICE)
+            
+            cached_latents = []
+            
+            with torch.no_grad():
+                for idx, img_path in enumerate(image_paths):
+                    img = Image.open(img_path).convert("RGB")
+                    img = img.resize((int(resolution), int(resolution)), Image.LANCZOS)
+                    img_tensor = torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0
+                    img_tensor = img_tensor.unsqueeze(0).to(DEVICE, dtype=DTYPE)
+                    img_tensor = 2.0 * img_tensor - 1.0
+                    
+                    latent = pipe.vae.encode(img_tensor).latent_dist.sample()
+                    latent = latent * vae_scaling_factor
+                    cached_latents.append(latent.cpu())
+                    
+                    del img_tensor, latent, img
+                    
+                    if idx % 2 == 0:
+                        torch.cuda.empty_cache()
+                        progress(0.32 + 0.08 * (idx / len(image_paths)), 
+                                desc=f"Encoding image {idx+1}/{len(image_paths)}")
+            
+            # Free VAE
+            pipe.vae.to("cpu")
+            del pipe.vae
+            aggressive_cleanup()
+            print(f"[Train] After VAE encoding: {get_gpu_memory_info()}")
+
+            # ============================================
+            # PHASE 5: Setup Transformer with Training Adapter
+            # ============================================
+            progress(0.42, desc="Setting up transformer with training adapter...")
+            
+            # Download training adapter
+            try:
+                adapter_path = hf_hub_download(
+                    repo_id="ostris/zimage_turbo_training_adapter",
+                    filename="zimage_turbo_training_adapter_v1.safetensors",
+                    local_dir=tmpdir
+                )
+                print(f"[Train] Training adapter downloaded: {adapter_path}")
+            except Exception as e:
+                return None, f"❌ Could not download training adapter: {e}"
+
+            # Get transformer (still on CPU from pipeline)
+            transformer = pipe.transformer
+            
+            # Load adapter via pipe's load_lora_weights
+            from safetensors.torch import load_file, save_file
+            
+            try:
+                pipe.load_lora_weights(adapter_path, adapter_name="training_adapter")
+                print("[Train] Training adapter loaded via load_lora_weights")
+            except Exception as e:
+                print(f"[Train] Warning: Could not load adapter via load_lora_weights: {e}")
+                
+            # Configure our training LoRA
+            progress(0.45, desc="Configuring LoRA...")
+
+            lora_config = LoraConfig(
+                r=int(lora_rank),
+                lora_alpha=int(lora_rank),
+                init_lora_weights="gaussian",
+                target_modules=[
+                    "to_q", "to_k", "to_v", "to_out.0",
+                    "attn.to_q", "attn.to_k", "attn.to_v", "attn.to_out.0",
+                    "ff.net.0.proj", "ff.net.2",
+                    "proj_in", "proj_out",
+                ],
+                lora_dropout=0.0,
+            )
+
+            transformer = get_peft_model(transformer, lora_config)
+            trainable = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
+            total = sum(p.numel() for p in transformer.parameters())
+            print(f"[Train] Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")
+
+            # Move to GPU
+            transformer.to(DEVICE)
+            print(f"[Train] Transformer on GPU: {get_gpu_memory_info()}")
+
+            if hasattr(transformer, 'enable_gradient_checkpointing'):
+                transformer.enable_gradient_checkpointing()
+
+            # Free the rest of pipeline
+            del pipe
+            aggressive_cleanup()
+
+            # Optimizer & Scheduler
+            optimizer = torch.optim.AdamW(
+                [p for p in transformer.parameters() if p.requires_grad],
+                lr=learning_rate, weight_decay=0.01, betas=(0.9, 0.999), eps=1e-8
+            )
+
+            from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, LinearLR, SequentialLR
+            
+            warmup_steps = min(100, int(num_steps * 0.1))
+            warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=warmup_steps)
+            cosine_scheduler = CosineAnnealingWarmRestarts(
+                optimizer, T_0=max(1, int(num_steps - warmup_steps)), eta_min=learning_rate * 0.01
+            )
+            lr_scheduler = SequentialLR(
+                optimizer, [warmup_scheduler, cosine_scheduler], milestones=[warmup_steps]
+            )
+
+            progress(0.50, desc=f"Training with {len(cached_latents)} samples...")
+            
+            # ============================================
+            # PHASE 6: Training Loop with CORRECT forward signature
+            # ZImageTransformer2DModel.forward(x, t, cap_feats, ...)
+            # x: List[Tensor] where each tensor is [C, F, H, W]
+            # cap_feats: List[Tensor] where each tensor is [seq_len, dim]
+            # ============================================
+            transformer.train()
+            losses = []
+            successful_steps = 0
+            
+            for step in range(int(num_steps)):
+                try:
+                    idx = np.random.randint(0, len(cached_latents))
+                    
+                    # latents: [B, C, H, W] -> need [C, F, H, W] where F=1
+                    latents = cached_latents[idx].to(DEVICE, dtype=DTYPE)
+                    # Remove batch dim, add frame dim: [1, C, H, W] -> [C, 1, H, W]
+                    latents = latents.squeeze(0).unsqueeze(1)  # [C, 1, H, W]
+                    
+                    # text_embeddings: [B, seq_len, dim] -> [seq_len, dim]
+                    text_embeddings = cached_text_embeddings[idx].to(DEVICE, dtype=DTYPE)
+                    text_embeddings = text_embeddings.squeeze(0)  # [seq_len, dim]
+
+                    # Timestep for flow matching (0 to 1)
+                    timesteps = torch.rand(1, device=DEVICE, dtype=DTYPE)
+                    
+                    # Create noisy latents using flow matching interpolation
+                    noise = torch.randn_like(latents)
+                    t = timesteps.view(-1, 1, 1, 1)
+                    noisy_latents = (1 - t) * latents + t * noise
+                    
+                    # Target is the velocity: noise - clean
+                    target = noise - latents
+
+                    # Scale timestep for model
+                    t_input = timesteps * 1000
+
+                    # CORRECT FORWARD CALL:
+                    # x and cap_feats must be Lists!
+                    with torch.amp.autocast('cuda', dtype=DTYPE):
+                        output = transformer(
+                            x=[noisy_latents],           # List of [C, F, H, W]
+                            t=t_input,                   # timestep
+                            cap_feats=[text_embeddings], # List of [seq_len, dim]
+                            return_dict=True
+                        )
+                    
+                    # Get model output - it will also be a list
+                    if hasattr(output, 'sample'):
+                        model_output = output.sample
+                        if isinstance(model_output, list):
+                            model_output = model_output[0]
+                    elif isinstance(output, tuple):
+                        model_output = output[0]
+                        if isinstance(model_output, list):
+                            model_output = model_output[0]
+                    else:
+                        model_output = output
+                        if isinstance(model_output, list):
+                            model_output = model_output[0]
+
+                    loss = compute_flow_matching_loss(model_output, target, timesteps)
+
+                    optimizer.zero_grad()
+                    loss.backward()
+                    torch.nn.utils.clip_grad_norm_(transformer.parameters(), 1.0)
+                    optimizer.step()
+                    lr_scheduler.step()
+
+                    losses.append(loss.item())
+                    successful_steps += 1
+                    
+                    del latents, text_embeddings, noise, noisy_latents, target, model_output, loss, output
+                    
+                    if step % 25 == 0:
+                        avg_loss = np.mean(losses[-50:]) if len(losses) >= 50 else np.mean(losses) if losses else float('nan')
+                        progress(
+                            0.50 + 0.40 * (step / int(num_steps)), 
+                            desc=f"Step {step}/{int(num_steps)} | Loss: {avg_loss:.4f}"
+                        )
+                        print(f"[Train] Step {step}: Loss={avg_loss:.4f}")
+
+                    if step % 100 == 0:
+                        gc.collect()
+                        torch.cuda.empty_cache()
+                        
+                except Exception as e:
+                    if step < 5:
+                        print(f"[Train] Error at step {step}: {e}")
+                        import traceback
+                        traceback.print_exc()
+                    continue
+
+            if successful_steps == 0:
+                return None, "❌ Training failed - no successful steps. Check model forward signature."
+
+            # ============================================
+            # PHASE 7: Save LoRA
+            # ============================================
+            progress(0.92, desc="Saving LoRA...")
+            
+            del cached_latents, cached_text_embeddings
+            aggressive_cleanup()
+
+            lora_state_dict = {}
+            for name, param in transformer.named_parameters():
+                if "lora" in name.lower() and param.requires_grad:
+                    clean_name = name.replace("base_model.model.", "")
+                    lora_state_dict[clean_name] = param.detach().cpu()
+
+            if not lora_state_dict:
+                return None, "❌ No LoRA weights found"
+
+            final_output = f"/tmp/{output_name}.safetensors"
+            save_file(lora_state_dict, final_output)
+            
+            file_size = os.path.getsize(final_output) / (1024 * 1024)
+            avg_final_loss = np.mean(losses[-100:]) if len(losses) >= 100 else np.mean(losses) if losses else float('nan')
+
+            training_info = f"""
+- Images: {len(image_paths)}
+- Steps: {successful_steps}
+- Final Loss: {avg_final_loss:.4f}
+- LR: {learning_rate}, Rank: {int(lora_rank)}, Resolution: {int(resolution)}
+"""
+
+            hub_result = ""
+            if upload_to_hub_flag:
+                progress(0.94, desc="Uploading to Hub...")
+                success, result = upload_to_hub(
+                    final_output, hub_repo_name or output_name, trigger_word, training_info
+                )
+                hub_result = f"\n\n🚀 Uploaded: {result}" if success else f"\n\n⚠️ Upload failed: {result}"
+
+            del transformer
+            aggressive_cleanup()
+            progress(1.0, desc="Complete!")
+
+            sample_captions = "\n".join([f"  - {c[:80]}..." for c in captions[:3]])
+
+            return final_output, f"""✅ Training complete!
+
+📁 LoRA: {output_name}.safetensors ({file_size:.1f} MB)
+🏷️ Trigger: {trigger_word}
+📊 Loss: {avg_final_loss:.4f}
+🖼️ Images: {len(image_paths)}
+⚙️ Steps: {successful_steps}
+
+**Sample captions:**
+{sample_captions}{hub_result}
+
+**Usage:**
+```python
+from diffusers import ZImagePipeline
+import torch
+
+pipe = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights("{output_name}.safetensors")
+pipe.to("cuda")
+
+image = pipe("{trigger_word}, your prompt here", num_inference_steps=8, guidance_scale=0.0).images[0]
+```"""
+
         except Exception as e:
-            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-supported LoRA")
-            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-supported LoRA"), gr.update(visible=False), gr.update(), "", None, ""
-    else:
-        return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
-
-def remove_custom_lora():
-    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
-
-run_lora.zerogpu = True
-
-css = '''
-#gen_btn{height: 100%}
-#gen_column{align-self: stretch}
-#title{text-align: center}
-#title h1{font-size: 3em; display:inline-flex; align-items:center}
-#title img{width: 100px; margin-right: 0.5em}
-#gallery .grid-wrap{height: 10vh}
-#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
-.card_internal{display: flex;height: 100px;margin-top: .5em}
-.card_internal img{margin-right: 1em}
-.styler{--form-gap-width: 0px !important}
-#progress{height:30px}
-#progress .generating{display:none}
-.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
-.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
-'''
-
-with gr.Blocks(delete_cache=(60, 60)) as demo:
-    title = gr.HTML(
-        """<h1>Z Image Turbo LoRA DLC 🧪</h1>""",
-        elem_id="title",
-    )
-    selected_index = gr.State(None)
+            aggressive_cleanup()
+            import traceback
+            return None, f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
+
+
+# ============================================
+# Gradio UI with Comic Style
+# ============================================
+with gr.Blocks(css=COMIC_CSS, theme=gr.themes.Soft(), title="Z-IMAGE GEN/LORA") as demo:
+    
+    # HOME Button
+    gr.HTML("""
+    <div class="home-button-container">
+        <a href="https://www.ginigen.com" target="_blank" class="home-button">
+            🏠 HOME
+        </a>
+        <span class="url-display">🌐 www.ginigen.com</span>
+    </div>
+    """)
+    
+    # Header
+    gr.HTML("""
+    <div class="header-container">
+        <div class="header-title">🎨 Z-IMAGE GEN/LORA 🎨</div>
+        <div class="header-subtitle">Train custom LoRA for Z-Image Turbo with Florence-2 auto-captioning</div>
+        <div style="margin-top:12px">
+            <span class="stats-badge">🧠 Florence-2 Caption</span>
+            <span class="stats-badge">⚡ Memory Optimized</span>
+            <span class="stats-badge">🚀 Hub Upload</span>
+            <span class="stats-badge">🎯 Person/Style/Object</span>
+        </div>
+    </div>
+    """)
+
+    # Status Row
+    gr.HTML('<div class="info-box">📊 <b>System Status</b> - Check GPU and HuggingFace connection</div>')
+    
     with gr.Row():
-        with gr.Column(scale=3):
-            prompt = gr.Textbox(label="Enter Prompt", lines=1, placeholder="✦︎ Choose the LoRA and type the prompt (LoRA = None → Base Model = Active)")
-        with gr.Column(scale=1, elem_id="gen_column"):
-            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
+        with gr.Column(scale=1):
+            gpu_status = gr.Textbox(label="🖥️ GPU Status", value=check_gpu(), interactive=False)
+        with gr.Column(scale=1):
+            hf_status = gr.Textbox(label="🔑 HF Token", value=check_hf_token(), interactive=False)
+        refresh_btn = gr.Button("🔄 Refresh", size="sm")
+        refresh_btn.click(fn=lambda: (check_gpu(), check_hf_token()), outputs=[gpu_status, hf_status])
+
     with gr.Row():
-        with gr.Column():
-            selected_info = gr.Markdown("### No LoRA Selected (Base Model)")
-            gallery = gr.Gallery(
-                [(item["image"], item["title"]) for item in loras],
-                label="Z-Image LoRAs",
-                allow_preview=False,
-                columns=3,
-                elem_id="gallery",
+        # Left Column - Images
+        with gr.Column(scale=1):
+            gr.HTML('<div class="info-box">📸 <b>Training Images</b> - Upload 6-20 high-quality images</div>')
+            images = gr.Gallery(
+                label="Upload Images", 
+                columns=4, 
+                height=300, 
+                type="filepath", 
+                interactive=True
             )
-            with gr.Group():
-                custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="Paste the LoRA path and press Enter (e.g., Shakker-Labs/AWPortrait-Z).")
-                gr.Markdown("[Check the list of Z-Image LoRA's](https://huggingface.co/models?other=base_model:adapter:Tongyi-MAI/Z-Image-Turbo)", elem_id="lora_list")
-            custom_lora_info = gr.HTML(visible=False)
-            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
-        with gr.Column():
-            progress_bar = gr.Markdown(elem_id="progress",visible=False)
-            result = gr.Image(label="Generated Image", format="png", height=630)
+            gr.HTML('<div class="tips-box">💡 <b>Tips:</b> 6-20 images, varied poses/angles, consistent subject, good lighting</div>')
 
-    with gr.Row():
-        with gr.Accordion("Advanced Settings", open=False):
+        # Right Column - Settings
+        with gr.Column(scale=1):
+            gr.HTML('<div class="info-box">⚙️ <b>Training Settings</b> - Configure your LoRA training</div>')
+            
+            trigger_word = gr.Textbox(
+                label="🏷️ Trigger Word", 
+                placeholder="ohwx_person", 
+                info="Use unique token like 'ohwx' to avoid conflicts"
+            )
+            output_name = gr.Textbox(
+                label="📁 Output Name", 
+                placeholder="my_lora"
+            )
+            
             with gr.Row():
-                input_image = gr.Image(label="Input image (Ignored for Z-Image-Turbo)", type="filepath", visible=False)
-                image_strength = gr.Slider(label="Denoise Strength", info="Ignored for Z-Image-Turbo", minimum=0.1, maximum=1.0, step=0.01, value=0.75, visible=False)
-            with gr.Column():
-                with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG Scale", info="Forced to 0.0 for Turbo", minimum=0, maximum=20, step=0.5, value=0.0, interactive=False)
-                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=25)
-                
-                with gr.Row():
-                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1536)
-                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1536)
-                
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
-
-    gallery.select(
-        update_selection,
-        inputs=[width, height],
-        outputs=[prompt, selected_info, selected_index, width, height]
-    )
-    custom_lora.input(
-        add_custom_lora,
-        inputs=[custom_lora],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
-    )
-    custom_lora_button.click(
-        remove_custom_lora,
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
-    )
-    gr.on(
-        triggers=[generate_button.click, prompt.submit],
-        fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
-        outputs=[result, seed, progress_bar]
+                use_auto_caption = gr.Checkbox(label="🔤 Auto-Caption (Florence-2)", value=True)
+                is_person_training = gr.Checkbox(label="👤 Person/Face Training", value=True)
+
+            with gr.Row():
+                num_steps = gr.Slider(500, 5000, 1500, step=100, label="🔢 Steps")
+                learning_rate = gr.Slider(1e-5, 5e-4, 5e-5, step=1e-5, label="📈 Learning Rate")
+
+            with gr.Row():
+                lora_rank = gr.Slider(4, 64, 32, step=4, label="🎚️ LoRA Rank")
+                resolution = gr.Slider(512, 1024, 1024, step=128, label="📐 Resolution")
+
+            batch_size = gr.Slider(1, 4, 1, step=1, visible=False)
+
+            gr.HTML('<div class="info-box">🚀 <b>Hub Upload</b> - Upload trained LoRA to HuggingFace</div>')
+            
+            with gr.Row():
+                upload_to_hub_flag = gr.Checkbox(label="📤 Upload to HF Hub (Private)", value=False)
+                hub_repo_name = gr.Textbox(label="📦 Repo Name", placeholder="my-zimage-lora")
+
+    # Train Button
+    with gr.Row():
+        train_btn = gr.Button("🚀 START TRAINING!", variant="primary", size="lg")
+
+    # Output
+    with gr.Row():
+        with gr.Column(scale=1):
+            output_file = gr.File(label="📥 Download LoRA")
+        with gr.Column(scale=1):
+            output_log = gr.Textbox(label="📋 Training Log", lines=15)
+
+    train_btn.click(
+        fn=train_lora,
+        inputs=[
+            images, trigger_word, output_name, num_steps, learning_rate, lora_rank,
+            resolution, batch_size, upload_to_hub_flag, hub_repo_name,
+            use_auto_caption, is_person_training
+        ],
+        outputs=[output_file, output_log]
     )
 
-demo.queue()
-demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)
+    # Recommended Settings Table
+    gr.HTML("""
+    <div class="info-box">
+        📋 <b>Recommended Settings by Use Case</b>
+        <table style="width:100%; margin-top:10px; border-collapse: collapse;">
+            <tr style="background:#3B82F6; color:white;">
+                <th style="padding:8px; border:2px solid #1F2937;">Use Case</th>
+                <th style="padding:8px; border:2px solid #1F2937;">Steps</th>
+                <th style="padding:8px; border:2px solid #1F2937;">LR</th>
+                <th style="padding:8px; border:2px solid #1F2937;">Rank</th>
+            </tr>
+            <tr style="background:#FEF9C3;">
+                <td style="padding:8px; border:2px solid #1F2937;">👤 Person</td>
+                <td style="padding:8px; border:2px solid #1F2937;">1500</td>
+                <td style="padding:8px; border:2px solid #1F2937;">5e-5</td>
+                <td style="padding:8px; border:2px solid #1F2937;">32</td>
+            </tr>
+            <tr style="background:#FFF;">
+                <td style="padding:8px; border:2px solid #1F2937;">🎨 Style</td>
+                <td style="padding:8px; border:2px solid #1F2937;">2000</td>
+                <td style="padding:8px; border:2px solid #1F2937;">1e-4</td>
+                <td style="padding:8px; border:2px solid #1F2937;">16</td>
+            </tr>
+            <tr style="background:#FEF9C3;">
+                <td style="padding:8px; border:2px solid #1F2937;">📦 Object</td>
+                <td style="padding:8px; border:2px solid #1F2937;">1200</td>
+                <td style="padding:8px; border:2px solid #1F2937;">8e-5</td>
+                <td style="padding:8px; border:2px solid #1F2937;">24</td>
+            </tr>
+        </table>
+    </div>
+    """)
+
+    # Footer
+    gr.HTML("""
+    <div class="footer-comic">
+        <p style="font-family:'Bangers',cursive;font-size:1.5rem;letter-spacing:2px">🎨 Z-IMAGE GEN/LORA 🎨</p>
+        <p>Powered by Z-Image Turbo + Florence-2 + PEFT</p>
+        <p>🧠 Auto Caption • ⚡ Memory Optimized • 🚀 Fast Training • 📤 Hub Upload</p>
+        <p style="margin-top:10px"><a href="https://www.ginigen.com" target="_blank" style="color:#FACC15;text-decoration:none;font-weight:bold;">🏠 www.ginigen.com</a></p>
+    </div>
+    """)
+
+if __name__ == "__main__":
+    mp.set_start_method('spawn', force=True)
+    demo.launch()