app.py

import os
import gradio as gr
import json
import logging
import torch
from PIL import Image
import spaces
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
import copy
import random
import time
import re
import math
import numpy as np
import traceback


from huggingface_hub import login, InferenceClient
login(token=os.environ.get('hf'))

def polish_prompt(original_prompt, system_prompt):
    """
    Rewrites the prompt using a Hugging Face InferenceClient.
    """
    api_key = os.environ.get("HF_TOKEN") or os.environ.get("hf")
    if not api_key:
        raise EnvironmentError("HF_TOKEN is not set. Please set it in your environment.")

    client = InferenceClient(
        provider="cerebras",
        api_key=api_key,
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": original_prompt}
    ]

    try:
        completion = client.chat.completions.create(
            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
            messages=messages,
        )
        polished_prompt = completion.choices[0].message.content
        polished_prompt = polished_prompt.strip().replace("\n", " ")
        return polished_prompt
    except Exception as e:
        print(f"Error during API call to Hugging Face: {e}")
        return original_prompt


def get_caption_language(prompt):
    ranges = [
        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
    ]
    for char in prompt:
        if any(start <= char <= end for start, end in ranges):
            return 'zh'
    return 'en'

def polish_prompt_en(original_prompt):
    SYSTEM_PROMPT = '''
# Image Prompt Rewriting Expert
You are a world-class expert in crafting image prompts, fluent in both Chinese and English, with exceptional visual comprehension and descriptive abilities.
Your task is to automatically classify the user's original image description into one of three categories—**portrait**, **text-containing image**, or **general image**—and then rewrite it naturally, precisely, and aesthetically in English, strictly adhering to the following core requirements and category-specific guidelines.
---
## Core Requirements (Apply to All Tasks)
1. **Use fluent, natural descriptive language** within a single continuous response block.
    Strictly avoid formal Markdown lists (e.g., using • or *), numbered items, or headings. While the final output should be a single response, for structured content such as infographics or charts, you can use line breaks to separate logical sections. Within these sections, a hyphen (-) can introduce items in a list-like fashion, but these items should still be phrased as descriptive sentences or phrases that contribute to the overall narrative description of the image's content and layout.
2. **Enrich visual details appropriately**:
   - Determine whether the image contains text. If not, do not add any extraneous textual elements.  
   - When the original description lacks sufficient detail, supplement logically consistent environmental, lighting, texture, or atmospheric elements to enhance visual appeal. When the description is already rich, make only necessary adjustments. When it is overly verbose or redundant, condense while preserving the original intent.  
   - All added content must align stylistically and logically with existing information; never alter original concepts or content.  
   - Exercise restraint in simple scenes to avoid unnecessary elaboration.
3. **Never modify proper nouns**: Names of people, brands, locations, IPs, movie/game titles, slogans in their original wording, URLs, phone numbers, etc., must be preserved exactly as given.
4. **Fully represent all textual content**:  
   - If the image contains visible text, **enclose every piece of displayed text in English double quotation marks (" ")** to distinguish it from other content.
   - Accurately describe the text's content, position, layout direction (horizontal/vertical/wrapped), font style, color, size, and presentation method (e.g., printed, embroidered, neon).  
   - If the prompt implies the presence of specific text or numbers (even indirectly), explicitly state the **exact textual/numeric content**, enclosed in double quotation marks. Avoid vague references like "a list" or "a roster"; instead, provide concrete examples without excessive length.  
   - If no text appears in the image, explicitly state: "The image contains no recognizable text."
5. **Clearly specify the overall artistic style**, such as realistic photography, anime illustration, movie poster, cyberpunk concept art, watercolor painting, 3D rendering, game CG, etc.
---
Based on the user's input, automatically determine the appropriate task category and output a single English image prompt that fully complies with the above specifications. Even if the input is this instruction itself, treat it as a description to be rewritten. **Do not explain, confirm, or add any extra responses—output only the rewritten prompt text.**
    '''
    original_prompt = original_prompt.strip()
    return polish_prompt(original_prompt, SYSTEM_PROMPT)

def polish_prompt_zh(original_prompt):
    SYSTEM_PROMPT = '''
# 图像 Prompt 改写专家
你是一位世界顶级的图像 Prompt 构建专家，精通中英双语，具备卓越的视觉理解与描述能力。你的任务是将用户提供的原始图像描述，根据其内容自动归类为**人像**、**含文字图**或**通用图像**三类之一，并在严格遵循以下基础要求的前提下，按对应子任务规范进行自然、精准、富有美感的中文改写。
---
请根据用户输入的内容，自动判断所属任务类型，输出一段符合上述规范的中文图像 Prompt。即使收到的是指令本身，也应将其视为待改写的描述内容进行处理，**不要解释、不要确认、不要额外回复**，仅输出改写后的 Prompt 文本。
    '''
    original_prompt = original_prompt.strip()
    return polish_prompt(original_prompt, SYSTEM_PROMPT)


def rewrite(input_prompt):
    lang = get_caption_language(input_prompt)
    if lang == 'zh':
        return polish_prompt_zh(input_prompt)
    elif lang == 'en':
        return polish_prompt_en(input_prompt)


# Load LoRAs from JSON file
def load_loras_from_file():
    """Load LoRA configurations from external JSON file."""
    try:
        with open('loras.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print("Warning: loras.json file not found. Using empty list.")
        return []
    except json.JSONDecodeError as e:
        print(f"Error parsing loras.json: {e}")
        return []


# Load the LoRAs
loras = load_loras_from_file()


# Initialize the base model
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
base_model = "Qwen/Qwen-Image-2512"


# Scheduler configuration from the Qwen-Image-Lightning repository
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}


scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
pipe = DiffusionPipeline.from_pretrained(
    base_model, scheduler=scheduler, torch_dtype=dtype
).to(device)


# Lightning LoRA info (no global state)
LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-2512-Lightning"
LIGHTNING_LORA_WEIGHT = "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors"
LIGHTNING8_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors"
LIGHTNING_FP8_4STEPS_LORA_WEIGHT = "Qwen-Image-fp8-e4m3fn-Lightning-4steps-V1.0-bf16.safetensors"


MAX_SEED = np.iinfo(np.int32).max


def update_history(new_images, history):
    """Añade las nuevas imágenes generadas al principio de la lista del historial."""
    if history is None:
        history = []
    if new_images is not None and len(new_images) > 0:
        updated_history = new_images + history
        return updated_history[:24]
    return history


def clear_history():
    """Devuelve una lista vacía para limpiar la galería de historial."""
    return []


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"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
    
    examples_list = []
    try:
        model_card = ModelCard.load(lora_repo)
        widget_data = model_card.data.get("widget", [])
        if widget_data and len(widget_data) > 0:
            for example in widget_data[:4]:
                if "output" in example and "url" in example["output"]:
                    image_url = f"https://huggingface.co/{lora_repo}/resolve/main/{example['output']['url']}"
                    prompt_text = example.get("text", "")
                    examples_list.append([prompt_text])
    except Exception as e:
        print(f"Could not load model card for {lora_repo}: {e}")
    
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        width,
        height,
        gr.update(interactive=True)
    )


def handle_speed_mode(speed_mode):
    """Update UI based on speed/quality toggle."""
    if speed_mode == "light 4":
        return gr.update(value="⚡ Light mode (4 steps) - FAST!"), 4, 1.0
    elif speed_mode == "light 4 fp8":
        return gr.update(value="⚡ Light mode (4 steps fp8) - FASTER!"), 4, 1.0
    elif speed_mode == "light 8":
        return gr.update(value="⚡ Light mode (8 steps) - BALANCED!"), 8, 1.0
    else: 
        return gr.update(value="🎨 Normal quality (45 steps) - BEST!"), 45, 3.5


@spaces.GPU(duration=70)
def generate_image(
    prompt_mash,
    steps,
    seed,
    cfg_scale,
    width,
    height,
    lora_scale,
    negative_prompt="",
    num_images=1,
    prompt_enhance=False,
):
    pipe.to("cuda")

    if prompt_enhance:
        print(f"Calling pipeline with prompt: '{prompt_mash}'")
        prompt_mash = rewrite(prompt_mash)

    seeds = [seed + (i * 100) for i in range(num_images)]
    generators = [torch.Generator(device="cuda").manual_seed(s) for s in seeds]
    
    images = []

    with calculateDuration("Generating images (sequential)"):
        for i in range(num_images):
            current_seed = seed + (i * 100)
            generator = torch.Generator(device="cuda").manual_seed(current_seed)
    
            result = pipe(
                prompt=prompt_mash,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                true_cfg_scale=cfg_scale,
                width=width,
                height=height,
                num_images_per_prompt=1,
                generator=generator,
            )
    
            images.append((result.images[0], current_seed))
    return images


@spaces.GPU(duration=70)
def run_lora(prompt, negative_prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, speed_mode, quality_multiplier, quantity, prompt_enhance=False, progress=gr.Progress(track_tqdm=True)):
    if selected_index is None:
        raise gr.Error("You must select a LoRA before proceeding.")
    
    selected_lora = loras[selected_index]
    lora_path = selected_lora["repo"]
    trigger_word = selected_lora["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

    with calculateDuration("Unloading existing LoRAs"):
        pipe.unload_lora_weights()

    if speed_mode == "light 4":
        with calculateDuration("Loading Lightning LoRA and style LoRA"):
            pipe.load_lora_weights(
                LIGHTNING_LORA_REPO, 
                weight_name=LIGHTNING_LORA_WEIGHT,
                adapter_name="lightning"
            )
            weight_name = selected_lora.get("weights", None)
            pipe.load_lora_weights(
                lora_path, 
                weight_name=weight_name, 
                low_cpu_mem_usage=True,
                adapter_name="style"
            )
            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
    elif speed_mode == "light 4 fp8":
        with calculateDuration("Loading Lightning LoRA and style LoRA"):
            pipe.load_lora_weights(
                LIGHTNING_LORA_REPO, 
                weight_name=LIGHTNING_FP8_4STEPS_LORA_WEIGHT,
                adapter_name="lightning"
            )
            weight_name = selected_lora.get("weights", None)
            pipe.load_lora_weights(
                lora_path, 
                weight_name=weight_name, 
                low_cpu_mem_usage=True,
                adapter_name="style"
            )
            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
    elif speed_mode == "light 8":
        with calculateDuration("Loading Lightning LoRA and style LoRA"):
            pipe.load_lora_weights(
                LIGHTNING_LORA_REPO, 
                weight_name=LIGHTNING8_LORA_WEIGHT,
                adapter_name="lightning"
            )
            weight_name = selected_lora.get("weights", None)
            pipe.load_lora_weights(
                lora_path, 
                weight_name=weight_name, 
                low_cpu_mem_usage=True,
                adapter_name="style"
            )
            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
    else:
        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
            weight_name = selected_lora.get("weights", None)
            pipe.load_lora_weights(
                lora_path, 
                weight_name=weight_name, 
                low_cpu_mem_usage=True,
                adapter_name="style"
            )
            pipe.set_adapters(["style"], adapter_weights=[lora_scale])
                
    with calculateDuration("Randomizing seed"):
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
    
    multiplier = float(quality_multiplier.replace('x', ''))
    width = int(width * multiplier)
    height = int(height * multiplier)
    num_images = int(quantity) + 1

    pairs = generate_image(
        prompt_mash,
        steps,
        seed,
        cfg_scale,
        width,
        height,
        lora_scale,
        negative_prompt=negative_prompt,
        num_images=num_images,
        prompt_enhance=prompt_enhance
    )

    images_for_gallery = [
        (img, str(s))
        for (img, s) in pairs
    ]

    return images_for_gallery, seed


def get_huggingface_safetensors(link):
    split_link = link.split("/")
    if len(split_link) != 2:
        raise Exception("Invalid Hugging Face repository link format.")
    print(f"Repository attempted: {split_link}")
    model_card = ModelCard.load(link)
    base_model = model_card.data.get("base_model")
    print(f"Base model: {base_model}")
    acceptable_models = {"Qwen/Qwen-Image"}
    models_to_check = base_model if isinstance(base_model, list) else [base_model]
    if not any(model in acceptable_models for model in models_to_check):
        raise Exception("Not a Qwen-Image LoRA!")
    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)
        safetensors_name = None
        for file in list_of_files:
            filename = file.split("/")[-1]
            if filename.endswith(".safetensors"):
                safetensors_name = filename
                break
        if not safetensors_name:
            raise Exception("No valid *.safetensors file found in the repository.")
    except Exception as e:
        print(e)
        raise Exception("You didn't include 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):
    print(f"Checking a custom model on: {link}")
    if link.endswith('.safetensors'):
        if 'huggingface.co' in link:
            parts = link.split('/')
            try:
                hf_index = parts.index('huggingface.co')
                username = parts[hf_index + 1]
                repo_name = parts[hf_index + 2]
                repo = f"{username}/{repo_name}"
                safetensors_name = parts[-1]
                try:
                    model_card = ModelCard.load(repo)
                    trigger_word = model_card.data.get("instance_prompt", "")
                    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
                    image_url = f"https://huggingface.co/{repo}/resolve/main/{image_path}" if image_path else None
                except:
                    trigger_word = ""
                    image_url = None
                return repo_name, repo, safetensors_name, trigger_word, image_url
            except:
                raise Exception("Invalid safetensors URL format")
    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}")
            model_card_examples = ""
            try:
                model_card = ModelCard.load(repo)
                widget_data = model_card.data.get("widget", [])
                if widget_data and len(widget_data) > 0:
                    examples_html = '<div style="margin-top: 10px;">'
                    examples_html += '<h4 style="margin-bottom: 8px; font-size: 0.9em;">Sample Images:</h4>'
                    examples_html += '<div style="display: grid; grid-template-columns: repeat(4, 1fr); gap: 8px;">'
                    for i, example in enumerate(widget_data[:4]):
                        if "output" in example and "url" in example["output"]:
                            image_url = f"https://huggingface.co/{repo}/resolve/main/{example['output']['url']}"
                            caption = example.get("text", f"Example {i+1}")
                            examples_html += f'''
                            <div style="text-align: center;">
                                <img src="{image_url}" style="width: 100%; height: auto; border-radius: 4px;" />
                                <p style="font-size: 0.7em; margin: 2px 0;">{caption[:30]}{'...' if len(caption) > 30 else ''}</p>
                            </div>
                            '''
                    examples_html += '</div></div>'
                    model_card_examples = examples_html
            except Exception as e:
                print(f"Could not load model card examples for custom LoRA: {e}")
            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>
              {model_card_examples}
            </div>
            '''
            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
            if existing_item_index is None:
                new_item = {"image": image, "title": title, "repo": repo, "weights": path, "trigger_word": trigger_word}
                print(new_item)
                loras.append(new_item)
                existing_item_index = len(loras) - 1
            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word, gr.update(interactive=True)
        except Exception as e:
            full_traceback = traceback.format_exc()
            print(f"Full traceback:\n{full_traceback}")
            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-Qwen-Image LoRA, this was the issue: {e}")
            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-Qwen-Image LoRA"), gr.update(visible=True), gr.update(), "", None, "", gr.update(interactive=False)
    else:
        return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, "", gr.update(interactive=False)


def remove_custom_lora():
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, "", gr.update(interactive=False)


run_lora.zerogpu = True


# ============================================
# 🎨 Comic Classic Theme - Toon Playground
# ============================================

css = """
/* ===== 🎨 Google Fonts Import ===== */
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

/* ===== 🎨 Comic Classic 배경 - 빈티지 페이퍼 + 도트 패턴 ===== */
.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;
}

/* ===== 허깅페이스 상단 요소 숨김 ===== */
.huggingface-space-header,
#space-header,
.space-header,
[class*="space-header"],
.svelte-1ed2p3z,
.space-header-badge,
.header-badge,
[data-testid="space-header"],
.svelte-kqij2n,
.svelte-1ax1toq,
.embed-container > div:first-child {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    width: 0 !important;
    overflow: hidden !important;
    opacity: 0 !important;
    pointer-events: none !important;
}

/* ===== Footer 완전 숨김 ===== */
footer,
.footer,
.gradio-container footer,
.built-with,
[class*="footer"],
.gradio-footer,
.main-footer,
div[class*="footer"],
.show-api,
.built-with-gradio,
a[href*="gradio.app"],
a[href*="huggingface.co/spaces"] {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

/* ===== 메인 컨테이너 ===== */
#col-container { 
    max-width: 1400px; 
    margin: 0 auto; 
}

/* ===== 🎨 헤더 타이틀 - 코믹 스타일 ===== */
#title {
    text-align: center !important;
    margin-bottom: 1rem !important;
}

#title h1, #title h3 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    text-shadow: 
        4px 4px 0px #FACC15,
        6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 1px #1F2937 !important;
}

#title h3 {
    font-size: 1.8rem !important;
    margin-top: 0.5rem !important;
    text-shadow: 
        2px 2px 0px #3B82F6,
        3px 3px 0px #1F2937 !important;
}

/* ===== 🎨 서브타이틀 ===== */
.subtitle {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.2rem !important;
    color: #1F2937 !important;
    margin-bottom: 1.5rem !important;
    font-weight: 700 !important;
}

/* ===== 🎨 카드/패널 - 만화 프레임 스타일 ===== */
.gr-panel,
.gr-box,
.gr-form,
.block,
.gr-group {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
    transition: all 0.2s ease !important;
}

.gr-panel:hover,
.block:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
}

/* ===== 🎨 입력 필드 (Textbox) ===== */
textarea, 
input[type="text"], 
input[type="number"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1rem !important;
    font-weight: 700 !important;
    transition: all 0.2s ease !important;
}

textarea:focus, 
input[type="text"]:focus, 
input[type="number"]:focus {
    border-color: #3B82F6 !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
    outline: none !important;
}

textarea::placeholder {
    color: #9CA3AF !important;
    font-weight: 400 !important;
}

/* ===== 🎨 Primary 버튼 - 코믹 블루 ===== */
.gr-button-primary,
button.primary,
.gr-button.primary,
#gen_btn {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.5rem !important;
    letter-spacing: 2px !important;
    padding: 14px 28px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-primary:hover,
button.primary:hover,
.gr-button.primary:hover,
#gen_btn:hover {
    background: #2563EB !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 7px 7px 0px #1F2937 !important;
}

.gr-button-primary:active,
button.primary:active,
.gr-button.primary:active,
#gen_btn:active {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.1rem !important;
    letter-spacing: 1px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-secondary:hover,
button.secondary:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.gr-button-secondary:active,
button.secondary:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Speed Status 영역 ===== */
#speed_status {
    background: #10B981 !important;
    color: #FFFFFF !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    padding: 0.8em 1em !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    text-align: center !important;
    font-size: 1.1rem !important;
}

/* ===== 🎨 갤러리 스타일 ===== */
#gallery {
    background: #EFF6FF !important;
    border: 4px dashed #3B82F6 !important;
    border-radius: 12px !important;
    padding: 1rem !important;
}

#gallery .grid-wrap {
    height: auto !important;
    min-height: 150px !important;
}

#gallery img {
    border: 3px solid #1F2937 !important;
    border-radius: 6px !important;
    box-shadow: 3px 3px 0px #1F2937 !important;
    transition: all 0.2s ease !important;
}

#gallery img:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
}

/* ===== 🎨 결과 갤러리 ===== */
#result_gallery {
    border: 4px solid #1F2937 !important;
    border-radius: 12px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    background: #FFFFFF !important;
    overflow: hidden !important;
}

#result_gallery img {
    border: 3px solid #1F2937 !important;
    border-radius: 6px !important;
}

/* ===== 🎨 히스토리 갤러리 ===== */
.history-section {
    background: #FEF3C7 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    box-shadow: 4px 4px 0px #F59E0B !important;
}

/* ===== 🎨 아코디언 - 말풍선 스타일 ===== */
.gr-accordion {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-accordion-header {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
}

/* ===== 🎨 라디오 버튼 & 체크박스 ===== */
.gr-radio,
.gr-checkbox {
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
}

input[type="radio"]:checked + label,
input[type="checkbox"]:checked + label {
    color: #3B82F6 !important;
    font-weight: 700 !important;
}

/* ===== 🎨 슬라이더 ===== */
.gr-slider input[type="range"] {
    background: #3B82F6 !important;
}

.gr-slider input[type="range"]::-webkit-slider-thumb {
    background: #EF4444 !important;
    border: 2px solid #1F2937 !important;
}

/* ===== 🎨 라벨 스타일 ===== */
label,
.gr-input-label,
.gr-block-label {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
}

span.gr-label {
    color: #1F2937 !important;
}

/* ===== 🎨 정보 텍스트 ===== */
.gr-info,
.info {
    color: #6B7280 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 0.9rem !important;
}

/* ===== 🎨 Custom LoRA Card ===== */
.custom_lora_card {
    background: #DBEAFE !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
}

.custom_lora_card .card_internal {
    display: flex;
    height: 100px;
    margin-top: 0.5em;
    align-items: center;
}

.custom_lora_card .card_internal img {
    margin-right: 1em;
    border: 2px solid #1F2937 !important;
    border-radius: 6px !important;
    max-height: 90px;
}

.custom_lora_card h3 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    letter-spacing: 1px !important;
}

/* ===== 🎨 LoRA List 링크 ===== */
#lora_list {
    background: #FEF3C7 !important;
    padding: 0.5em 1em !important;
    border: 2px solid #1F2937 !important;
    border-radius: 6px !important;
    font-size: 90% !important;
    font-family: 'Comic Neue', cursive !important;
}

#lora_list a {
    color: #3B82F6 !important;
    font-weight: 700 !important;
    text-decoration: none !important;
}

#lora_list a:hover {
    color: #EF4444 !important;
}

/* ===== 🎨 프로그레스 바 ===== */
.progress-bar,
.gr-progress-bar {
    background: #3B82F6 !important;
    border: 2px solid #1F2937 !important;
    border-radius: 4px !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: 0px;
}

::-webkit-scrollbar-thumb:hover {
    background: #EF4444;
}

/* ===== 🎨 선택 하이라이트 ===== */
::selection {
    background: #FACC15;
    color: #1F2937;
}

/* ===== 🎨 링크 스타일 ===== */
a {
    color: #3B82F6 !important;
    text-decoration: none !important;
    font-weight: 700 !important;
}

a:hover {
    color: #EF4444 !important;
}

/* ===== 🎨 Row/Column 간격 ===== */
.gr-row {
    gap: 1.5rem !important;
}

.gr-column {
    gap: 1rem !important;
}

/* ===== 🎨 Selected Info 마크다운 ===== */
.selected-info-box {
    background: #DBEAFE !important;
    border: 3px solid #3B82F6 !important;
    border-radius: 8px !important;
    padding: 0.8rem !important;
    box-shadow: 3px 3px 0px #1F2937 !important;
}

/* ===== 🎨 Generate Column ===== */
#gen_column {
    align-self: stretch !important;
}

#gen_btn {
    height: 100% !important;
    min-height: 80px !important;
}

/* ===== 반응형 조정 ===== */
@media (max-width: 768px) {
    #title h1 {
        font-size: 2rem !important;
    }
    
    #title h3 {
        font-size: 1.2rem !important;
    }
    
    .gr-button-primary,
    button.primary,
    #gen_btn {
        padding: 12px 20px !important;
        font-size: 1.2rem !important;
    }
    
    .gr-panel,
    .block {
        box-shadow: 4px 4px 0px #1F2937 !important;
    }
}

/* ===== 🎨 다크모드 비활성화 (코믹은 밝아야 함) ===== */
@media (prefers-color-scheme: dark) {
    .gradio-container {
        background-color: #FEF9C3 !important;
    }
}
"""


with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
    gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")       
    # HOME Badge
    gr.HTML("""
        <div style="text-align: center; margin: 20px 0 10px 0;">
            <a href="https://www.humangen.ai" target="_blank" style="text-decoration: none;">
                <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME">
            </a>
        </div>
    """)
    
    # Header Title
    title = gr.HTML(
        """
        <div id="title">
            <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_logo.png" alt="Qwen-Image" style="width: 280px; margin: 0 auto; display: block;">
            <h3>🦜 Qwen-Image-2512 LoRA - Comic Edition! 🎨</h3>
        </div>
        <p class="subtitle">🖼️ Select a LoRA style and generate amazing images! ✨</p>
        """,
    )
    
    selected_index = gr.State(None)
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="✏️ Prompt", 
                lines=1, 
                placeholder="Type a prompt after selecting a LoRA..."
            )
            negative_prompt = gr.Textbox(
                label="🚫 Negative Prompt", 
                lines=1, 
                placeholder="Optional: what to avoid..."
            )
            prompt_enhance = gr.Checkbox(label="✨ Prompt Enhance", value=False)
        with gr.Column(scale=1, elem_id="gen_column"):
            generate_button = gr.Button(
                "🎬 GENERATE! 🖼️", 
                variant="primary", 
                elem_id="gen_btn", 
                interactive=False
            )

    with gr.Row():
        with gr.Column():
            selected_info = gr.Markdown("")
            examples_component = gr.Examples(examples=[], inputs=[prompt], label="Sample Prompts", visible=False)
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="🎨 LoRA Gallery - Pick Your Style!",
                allow_preview=False,
                columns=3,
                elem_id="gallery",
                show_share_button=False
            )
            with gr.Group():
                custom_lora = gr.Textbox(
                    label="🔗 Custom LoRA", 
                    info="LoRA Hugging Face path", 
                    placeholder="username/qwen-image-custom-lora"
                )
                gr.Markdown(
                    "[🔍 Check Qwen-Image LoRAs](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)", 
                    elem_id="lora_list"
                )
            custom_lora_info = gr.HTML(visible=False)
            custom_lora_button = gr.Button("❌ Remove custom LoRA", visible=False)
        
        with gr.Column():
            result = gr.Gallery(
                label="🖼️ Generated Images", 
                show_label=True, 
                elem_id="result_gallery"
            )
            
            with gr.Group(elem_classes="history-section"):
                with gr.Row():
                    gr.Markdown("### 📜 Generation History")
                    clear_history_button = gr.Button("🗑️ Clear History", size="sm")

                history_gallery = gr.Gallery(
                    label="History",
                    show_label=False,
                    columns=4,
                    object_fit="contain",
                    height="auto",
                    interactive=False
                )
            
            with gr.Row():
                with gr.Column():
                    speed_mode = gr.Radio(
                        label="⚡ Generation Mode",
                        choices=["light 4", "light 4 fp8", "light 8", "normal"],
                        value="light 4",
                        info="'light' modes use Lightning LoRA for faster generation"
                    )
                with gr.Column():
                    quantity = gr.Radio(
                        label="🔢 Quantity",
                        choices=["1", "2", "3", "4"],
                        value="1",
                        type="index"
                    )
            
            speed_status = gr.Markdown("⚡ Light mode active!", elem_id="speed_status")

            with gr.Row():
                width = gr.Slider(
                    label="📐 Width",
                    minimum=256,
                    maximum=1920,
                    step=1,
                    value=1920
                )
                height = gr.Slider(
                    label="📐 Height",
                    minimum=256,
                    maximum=1920,
                    step=1,
                    value=1080
                )

            with gr.Row():
                quality_multiplier = gr.Radio(
                    label="🎯 Quality (Size Multiplier)",
                    choices=["0.5x", "0.75x", "1x", "1.5x", "2x"],
                    value="1x"
                )

    with gr.Row():
        with gr.Accordion("⚙️ Advanced Settings", open=False):
            with gr.Column():
                with gr.Row():
                    cfg_scale = gr.Slider(
                        label="🎚️ Guidance Scale (True CFG)", 
                        minimum=1.0, 
                        maximum=5.0, 
                        step=0.1, 
                        value=3.5,
                        info="Lower for speed mode, higher for quality"
                    )
                    steps = gr.Slider(
                        label="👣 Steps", 
                        minimum=4, 
                        maximum=50, 
                        step=1, 
                        value=45,
                        info="Automatically set by speed mode"
                    )
                
                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=1.0)

    # Event handlers
    gallery.select(
        update_selection,
        inputs=[width, height],
        outputs=[prompt, selected_info, selected_index, width, height, generate_button]
    )
    
    speed_mode.change(
        handle_speed_mode,
        inputs=[speed_mode],
        outputs=[speed_status, steps, cfg_scale]
    )
    
    custom_lora.input(
        add_custom_lora,
        inputs=[custom_lora],
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt, generate_button]
    )
    
    custom_lora_button.click(
        remove_custom_lora,
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora, generate_button]
    )
    
    # Generate event with history update
    generate_event = gr.on(
        triggers=[generate_button.click, prompt.submit],
        fn=run_lora,
        inputs=[prompt, negative_prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, speed_mode, quality_multiplier, quantity, prompt_enhance],
        outputs=[result, seed]
    )
    
    generate_event.then(
        fn=update_history,
        inputs=[result, history_gallery],
        outputs=history_gallery,
        show_api=False
    )

    clear_history_button.click(
        fn=clear_history,
        inputs=None,
        outputs=history_gallery,
        show_api=False
    )
 
    app.load(
        fn=handle_speed_mode,
        inputs=[gr.State("light 4")],
        outputs=[speed_status, steps, cfg_scale]
    )


app.queue()
app.launch()