"""
Qwen Image to LoRA Generator
🎨 Comic Classic Theme
⚡ Optimized with Lazy Loading for Fast Startup
"""

import gradio as gr
import numpy as np
import torch
import random
import spaces
from ulid import ULID
from safetensors.torch import save_file
from PIL import Image
import os

DTYPE = torch.bfloat16
MAX_SEED = np.iinfo(np.int32).max

# ============================================
# 🚀 LAZY LOADING - 모델을 처음 사용할 때만 로딩
# ============================================

pipe_lora = None
pipe_imagen = None
lora_model_loaded = False
imagen_model_loaded = False


def get_vram_config_disk_offload():
    return {
        "offload_dtype": "disk",
        "offload_device": "disk",
        "onload_dtype": "disk",
        "onload_device": "disk",
        "preparing_dtype": torch.bfloat16,
        "preparing_device": "cuda",
        "computation_dtype": torch.bfloat16,
        "computation_device": "cuda",
    }


def get_vram_config():
    return {
        "offload_dtype": "disk",
        "offload_device": "disk",
        "onload_dtype": torch.bfloat16,
        "onload_device": "cuda",
        "preparing_dtype": torch.bfloat16,
        "preparing_device": "cuda",
        "computation_dtype": torch.bfloat16,
        "computation_device": "cuda",
    }


# loras 디렉토리 생성
os.makedirs("loras", exist_ok=True)


@spaces.GPU(duration=300)
def generate_lora(input_images, progress=gr.Progress(track_tqdm=True)):
    """LoRA 생성 - GPU 컨텍스트 안에서 모델 로딩"""
    global pipe_lora, lora_model_loaded
    
    if not input_images:
        gr.Warning("⚠️ Please upload at least one image!")
        return None, gr.update(interactive=False), gr.update(interactive=False)
    
    # GPU 컨텍스트 안에서 모델 로딩
    if not lora_model_loaded:
        print("🔄 Loading LoRA generation pipeline...")
        
        from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig
        
        vram_config_disk_offload = get_vram_config_disk_offload()
        
        pipe_lora = QwenImagePipeline.from_pretrained(
            torch_dtype=torch.bfloat16,
            device="cuda",
            model_configs=[
                ModelConfig(
                    download_source="huggingface", 
                    model_id="DiffSynth-Studio/General-Image-Encoders", 
                    origin_file_pattern="SigLIP2-G384/model.safetensors", 
                    **vram_config_disk_offload
                ),
                ModelConfig(
                    download_source="huggingface", 
                    model_id="DiffSynth-Studio/General-Image-Encoders", 
                    origin_file_pattern="DINOv3-7B/model.safetensors", 
                    **vram_config_disk_offload
                ),
                ModelConfig(
                    download_source="huggingface", 
                    model_id="DiffSynth-Studio/Qwen-Image-i2L", 
                    origin_file_pattern="Qwen-Image-i2L-Style.safetensors", 
                    **vram_config_disk_offload
                ),
            ],
            processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"),
            vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
        )
        lora_model_loaded = True
        print("✅ LoRA pipeline loaded!")
    
    pipeline = pipe_lora
    
    from diffsynth.pipelines.qwen_image import (
        QwenImageUnit_Image2LoRAEncode, 
        QwenImageUnit_Image2LoRADecode
    )
    
    ulid = str(ULID()).lower()[:12]
    print(f"🎯 Generating LoRA with ID: {ulid}")
    
    input_images = [Image.open(filepath).convert("RGB") for filepath, _ in input_images]
    
    with torch.no_grad():
        embs = QwenImageUnit_Image2LoRAEncode().process(pipeline, image2lora_images=input_images)
        lora = QwenImageUnit_Image2LoRADecode().process(pipeline, **embs)["lora"]
    
    lora_name = f"{ulid}.safetensors"
    lora_path = f"loras/{lora_name}"
    
    save_file(lora, lora_path)
    print(f"✅ LoRA saved: {lora_path}")
    
    return lora_name, gr.update(interactive=True, value=lora_path), gr.update(interactive=True)


@spaces.GPU(duration=300)
def generate_image(
    lora_name, prompt, negative_prompt="blurry ugly bad",
    width=1024, height=1024, seed=42, randomize_seed=True,
    guidance_scale=3.5, num_inference_steps=8,
    progress=gr.Progress(track_tqdm=True)
):
    """이미지 생성 - GPU 컨텍스트 안에서 모델 로딩"""
    global pipe_imagen, imagen_model_loaded
    
    if not lora_name:
        gr.Warning("⚠️ Please generate a LoRA first!")
        return None, seed
    
    # GPU 컨텍스트 안에서 모델 로딩
    if not imagen_model_loaded:
        print("🔄 Loading Image generation pipeline...")
        
        from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig
        
        vram_config = get_vram_config()
        
        pipe_imagen = QwenImagePipeline.from_pretrained(
            torch_dtype=torch.bfloat16,
            device="cuda",
            model_configs=[
                ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config),
                ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", **vram_config),
                ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
            ],
            tokenizer_config=ModelConfig(download_source="huggingface", model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
            vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
        )
        imagen_model_loaded = True
        print("✅ Image generation pipeline loaded!")
    
    pipeline = pipe_imagen
    
    lora_path = f"loras/{lora_name}"
    
    if not os.path.exists(lora_path):
        gr.Warning(f"⚠️ LoRA file not found: {lora_path}")
        return None, seed
    
    pipeline.clear_lora()
    pipeline.load_lora(pipeline.dit, lora_path)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    print(f"🎨 Generating image with seed: {seed}")
    
    output_image = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
    )
    
    print("✅ Image generated!")
    
    return output_image, seed


# ============================================
# 🎨 Comic Classic Theme CSS
# ============================================

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: 960px; 
    margin: 0 auto; 
}

/* ===== 🎨 헤더 타이틀 - 코믹 스타일 ===== */
.header-text h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3.5rem !important;
    font-weight: 400 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 
        4px 4px 0px #FACC15,
        6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 2px #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;
}

/* ===== 🎨 입력 필드 ===== */
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 {
    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.3rem !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 {
    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 {
    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;
}

/* ===== 🎨 라벨 스타일 ===== */
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-file-upload {
    border: 3px dashed #1F2937 !important;
    border-radius: 8px !important;
    background: #FEF9C3 !important;
}

.gr-file-upload:hover {
    border-color: #3B82F6 !important;
    background: #EFF6FF !important;
}

/* ===== 🎨 갤러리 스타일 ===== */
.gr-gallery {
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

/* ===== 🎨 이미지 출력 영역 ===== */
.gr-image,
.image-container {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    overflow: hidden !important;
    background: #FFFFFF !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;
}

/* ===== 🎨 슬라이더 ===== */
input[type="range"] {
    accent-color: #EF4444 !important;
    height: 8px !important;
}

/* ===== 🎨 체크박스 ===== */
input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 20px !important;
    height: 20px !important;
    border: 2px solid #1F2937 !important;
}

/* ===== 🎨 다운로드 버튼 ===== */
.gr-download-button,
button[class*="download"] {
    background: #10B981 !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;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-download-button:hover,
button[class*="download"]:hover {
    background: #059669 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !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;
}

/* ===== 🎨 구분선 ===== */
hr {
    border: none !important;
    border-top: 3px dashed #1F2937 !important;
    margin: 1.5rem 0 !important;
}

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

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

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

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


# ============================================
# 🎮 Gradio UI
# ============================================

with gr.Blocks(title="Qwen Image to LoRA") as demo:
    # gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")       
    # CSS 주입
    gr.HTML(f"<style>{css}</style>")
    
    # 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
    gr.Markdown("# 🎨 QWEN IMAGE TO LORA 🎨", elem_classes="header-text")
    gr.Markdown('<p class="subtitle">✨ Generate LoRA from your images instantly! 🖼️</p>')
    
    with gr.Column(elem_id="col-container"):
        
        # ===== STEP 1: LoRA Generation =====
        gr.Markdown("### 📤 STEP 1: Upload Images & Generate LoRA")
        
        with gr.Row():
            with gr.Column():
                input_images = gr.Gallery(
                    label="📁 Input Images (Upload 1-5 images)", 
                    file_types=["image"],
                    columns=2, 
                    object_fit="cover", 
                    height=300
                )
                lora_button = gr.Button("🚀 Generate LoRA", variant="primary", size="lg")
            
            with gr.Column():
                lora_name = gr.Textbox(
                    label="📄 Generated LoRA Name", 
                    lines=1, 
                    interactive=False,
                    placeholder="LoRA name will appear here..."
                )
                lora_download = gr.DownloadButton(
                    label="📥 Download LoRA", 
                    interactive=False
                )
                gr.Markdown("""
                    <div style="background: #ECFDF5; border: 2px solid #10B981; border-radius: 8px; padding: 10px; margin-top: 10px;">
                        <p style="margin: 0; font-size: 0.9rem; color: #065F46;">
                            💡 <strong>Tip:</strong> Upload 1-5 images with similar style for best results!
                        </p>
                    </div>
                """)
        
        gr.Markdown("---")
        
        # ===== STEP 2: Image Generation =====
        gr.Markdown("### 🖼️ STEP 2: Generate Image with Your LoRA")
        
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="💬 Prompt",
                    lines=3,
                    placeholder="Describe the image you want to generate...",
                    value="a man in a fishing boat on a calm lake at sunset"
                )
                
                imagen_button = gr.Button(
                    "🎨 Generate Image", 
                    variant="primary", 
                    interactive=False,
                    size="lg"
                )
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    negative_prompt = gr.Textbox(
                        label="❌ Negative Prompt",
                        lines=2,
                        value="blurry, ugly, bad quality, distorted"
                    )
                    
                    num_inference_steps = gr.Slider(
                        label="🔄 Steps", 
                        minimum=1, 
                        maximum=50, 
                        step=1, 
                        value=25,
                        info="More steps = better quality but slower"
                    )
                    
                    with gr.Row():
                        width = gr.Slider(
                            label="📐 Width", 
                            minimum=512, 
                            maximum=1280, 
                            step=32, 
                            value=768
                        )
                        height = gr.Slider(
                            label="📐 Height", 
                            minimum=512, 
                            maximum=1280, 
                            step=32, 
                            value=1024
                        )
                    
                    with gr.Row():
                        seed = gr.Slider(
                            label="🎯 Seed", 
                            minimum=0, 
                            maximum=MAX_SEED, 
                            step=1, 
                            value=42
                        )
                        guidance_scale = gr.Slider(
                            label="🎚️ Guidance Scale", 
                            minimum=0.0, 
                            maximum=10.0, 
                            step=0.1, 
                            value=3.5
                        )
                    
                    randomize_seed = gr.Checkbox(
                        label="🎲 Randomize Seed", 
                        value=True
                    )
            
            with gr.Column():
                output_image = gr.Image(
                    label="🖼️ Generated Image",
                    height=500
                )
                
                used_seed = gr.Number(
                    label="🎯 Seed Used",
                    interactive=False
                )
    
    # ===== Event Handlers =====
    lora_button.click(
        fn=generate_lora,
        inputs=[input_images],
        outputs=[lora_name, lora_download, imagen_button]
    )
    
    imagen_button.click(
        fn=generate_image,
        inputs=[
            lora_name, 
            prompt, 
            negative_prompt, 
            width, 
            height, 
            seed, 
            randomize_seed, 
            guidance_scale, 
            num_inference_steps
        ],
        outputs=[output_image, used_seed]
    )


if __name__ == "__main__":
    demo.launch(ssr_mode=False, server_name="0.0.0.0", share=True)