app.py

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
# ============================================
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;
}

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;
}

.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;
}

/* Image/Gallery Container */
.gr-image, .gr-gallery {
    border: 3px solid #1F2937 !important;
    border-radius: 8px !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/VxF2jh087d4V6f9LWw1jQoLcQ_ymdhyD7XtRnq2KVYmN6DZsL4vCTrd1v-ubr8zfejyCCKvUWaBVf9JM9GqR271X6nh4e6mUbW11LaFr9QtepztFJeDZJ1VISkW5KBbebCpqv-w2Uv7RmMPwB5kacg.webp",
        
        "title": "koyoon Style",
        "repo": "./",  # 로컬 경로
        "weights": "koyoon.safetensors",
        "trigger_word": "koyoon"    
    },
    
    {
        "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/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}) ✅"
    
    # Default aspect ratio
    aspect = "1:1 (Instagram Square)"
    width = 1024
    height = 1024
    
    if "aspect" in selected_lora:
        if selected_lora["aspect"] == "portrait":
            aspect = "9:16 (Instagram Reels/TikTok/Shorts)"
            width = 768
            height = 1344
        elif selected_lora["aspect"] == "landscape":
            aspect = "16:9 (YouTube/Twitter/X)"
            width = 1344
            height = 768
    
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        aspect,
        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:
    gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")        
    # HOME Button
    gr.HTML("""
    <div class="home-button-container">
        <a href="https://www.humangen.ai" target="_blank" class="home-button">
            🏠 HOME
        </a>
        <span class="url-display">🌐 www.humangen.ai</span>
    </div>
    """)
    
    # Header
    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)

    # SNS Aspect Ratio Presets
    ASPECT_RATIOS = {
        "1:1 (Instagram Square)": (1024, 1024),
        "9:16 (Instagram Reels/TikTok/Shorts)": (768, 1344),
        "16:9 (YouTube/Twitter/X)": (1344, 768),
        "4:5 (Instagram Portrait)": (896, 1120),
        "5:4 (Instagram Landscape)": (1120, 896),
        "3:4 (Portrait Photo)": (896, 1152),
        "4:3 (Landscape Photo)": (1152, 896),
        "2:3 (Pinterest)": (832, 1248),
        "3:2 (Classic Photo)": (1248, 832),
        "21:9 (Cinematic Ultra-wide)": (1344, 576),
        "9:21 (Tall Banner)": (576, 1344),
    }
    
    def update_size(aspect_ratio):
        width, height = ASPECT_RATIOS.get(aspect_ratio, (1024, 1024))
        return width, height
    
    with gr.Row():
        with gr.Accordion("⚙️ Advanced Settings", open=True):
            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)
            
            gr.HTML('<div class="info-box">📐 <b>Image Size</b> - Select aspect ratio for different platforms</div>')
            
            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    choices=list(ASPECT_RATIOS.keys()),
                    value="1:1 (Instagram Square)",
                    label="📱 Aspect Ratio (SNS Presets)",
                    info="Choose the best ratio for your target platform"
                )
            
            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.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():
                    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)
    
    # Connect aspect ratio dropdown to width/height sliders
    aspect_ratio.change(
        fn=update_size,
        inputs=[aspect_ratio],
        outputs=[width, height]
    )

    # 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.humangen.ai" target="_blank" style="color:#FACC15;text-decoration:none;font-weight:bold;">🏠 www.humangen.ai</a></p>
    </div>
    """)

    gallery.select(
        update_selection,
        inputs=[width, height],
        outputs=[prompt, selected_info, selected_index, aspect_ratio, 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)