zlora

Running on Zero

App Files Files Community

seawolf2357 commited on 4 days ago

Commit

d32fd30

verified ·

1 Parent(s): 4a310f7

Update app.py

Browse files

Files changed (1) hide show

app.py +559 -877

app.py CHANGED Viewed

@@ -1,34 +1,32 @@
 import os
-import gc
-import torch
-import gradio as gr
-from PIL import Image
-from pathlib import Path
-import shutil
 import json
-from huggingface_hub import hf_hub_download, HfApi, create_repo, upload_file
-import tempfile
-import multiprocessing as mp
-import spaces  # ZeroGPU support
-# Training imports
-from peft import LoraConfig, get_peft_model
-from tqdm.auto import tqdm
 import numpy as np
-# 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
@@ -202,15 +200,6 @@ textarea:focus, input[type="text"]:focus {
     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;
@@ -224,19 +213,6 @@ label, .gr-input-label, .gr-block-label {
     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;
@@ -284,11 +260,6 @@ 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;
@@ -296,736 +267,495 @@ input[type="range"] {
     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;
-}
 """
-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()
-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!")
-        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()}")
-    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
-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 = []
-    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:
-            continue
-        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:
-        captions = [trigger_word] * len(image_paths)
-    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()
-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"
-    try:
-        api = HfApi(token=HF_TOKEN)
-        username = get_hf_username()
-        if not username:
-            return False, "Could not get username"
-        repo_id = f"{username}/{repo_name}"
-        try:
-            create_repo(repo_id=repo_id, token=HF_TOKEN, private=True, repo_type="model", exist_ok=True)
-        except:
-            pass
-        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)
-@spaces.GPU(duration=600)  # 10 minutes max for training
-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:
-            aggressive_cleanup()
-            import traceback
-            return None, f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
-# ============================================
-# Gradio UI with Comic Style
-# ============================================
-with gr.Blocks(title="Z-IMAGE GEN/LORA") as demo:
     # HOME Button
     gr.HTML("""
@@ -1041,138 +771,90 @@ with gr.Blocks(title="Z-IMAGE GEN/LORA") as demo:
     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=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():
-        # 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
             )
-            gr.HTML('<div class="tips-box">💡 <b>Tips:</b> 6-20 images, varied poses/angles, consistent subject, good lighting</div>')
-        # 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():
-                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]
-    )
-    # 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(css=COMIC_CSS, theme=gr.themes.Soft(), ssr_mode=False)

 import os
 import json
+import copy
+import time
+import requests
+import random
+import logging
 import numpy as np
+import spaces
+from typing import Any, Dict, List, Optional, Union
+import torch
+from PIL import Image
+import gradio as gr
+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
 # ============================================
 # Comic Style CSS
     color: #1F2937 !important;
 }
 label, .gr-input-label, .gr-block-label {
     color: #1F2937 !important;
     font-family: 'Comic Neue', cursive !important;
     box-shadow: 4px 4px 0 #1F2937 !important;
 }
 .footer-comic {
     text-align: center;
     padding: 20px;
     accent-color: #3B82F6;
 }
 /* Image/Gallery Container */
 .gr-image, .gr-gallery {
     border: 3px solid #1F2937 !important;
     box-shadow: 4px 4px 0 #1F2937 !important;
 }
+/* Original CSS additions */
+#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}
 """
+loras = [
+    # 로컬 jimin LoRA (app.py와 같은 디렉토리에 jimin.safetensors 필요)
+    {
+        "image": "https://i.namu.wiki/i/umL8EZtn0hs-nMRYeFxIrkGrMe-R1u5c9fJE8ufrLjvXz52VcSIbG7TT9QJoL2rR7vsFww1lLrE4bwfn5uOBzfq9a90HGdNdlTLmr_KoqOchTovbVC3RDzhDbp7FI-Wq-esCu7_BYIptqethL4onBg.webp",
+        "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
+    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,
     )
+@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()
+    # 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}"
         else:
+            prompt_mash = prompt
+        # 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.")
     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)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    # Note: Z-Image-Turbo is strictly T2I in this reference implementation.
+    # Img2Img via image_input is disabled/ignored for this pipeline update.
+    with calculateDuration("Generating image"):
+        # For Turbo models, guidance_scale is typically 0.0
+        forced_guidance = 0.0 # Turbo mode
+        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]
+    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)
+            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
         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
+with gr.Blocks(title="Z-IMAGE GEN/LORA", delete_cache=(60, 60)) as demo:
     # HOME Button
     gr.HTML("""
     gr.HTML("""
     <div class="header-container">
         <div class="header-title">🎨 Z-IMAGE GEN/LORA 🎨</div>
+        <div class="header-subtitle">Generate amazing images with Z-Image Turbo and various LoRA styles!</div>
         <div style="margin-top:12px">
+            <span class="stats-badge">⚡ Turbo Speed</span>
+            <span class="stats-badge">🎭 30+ LoRAs</span>
+            <span class="stats-badge">🖼️ High Quality</span>
+            <span class="stats-badge">🔧 Custom LoRA</span>
         </div>
     </div>
     """)
+    selected_index = gr.State(None)
     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.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",
             )
+            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)
     with gr.Row():
+        with gr.Accordion("Advanced Settings", open=False):
+            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)
     # 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 + LoRA Adapters</p>
+        <p>⚡ Fast Generation • 🎭 Multiple Styles • 🖼️ High Quality</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>
     """)
+    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]
+    )
+demo.queue()
+demo.launch(css=COMIC_CSS, ssr_mode=False, show_error=True)