print("Starting CoTyle app...")

import os
import sys
import json
import torch
import gradio as gr
import spaces
from PIL import Image
from huggingface_hub import snapshot_download
import gc
try:
    import pynvml
    pynvml.nvmlInit()
    device_count = pynvml.nvmlDeviceGetCount()
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        name = pynvml.nvmlDeviceGetName(handle)
        if isinstance(name, bytes):
            name = name.decode('utf-8')
        print(f"GPU {i}: {name}")
except Exception as e:
    print("无法获取 GPU 信息:", e)

REPO_ID = "Kwai-Kolors/cotyle"
HF_TOKEN = os.getenv("HF_TOKEN")

device = "cuda" if torch.cuda.is_available() else "cpu"
weight_type = torch.bfloat16 if device == "cuda" else torch.float32

SUGGESTED_PROMPTS = [
    "An artist sits outdoors, engrossed in their work, brush in hand, capturing the scene with focused intensity. On the canvas, trees and buildings blend seamlessly with the real-world surroundings. Symbols from different cultures, along with animals, plants, and abstract lines, float around them. As the brush touches the canvas, the paint transforms into points of light that scatter, while sheets of paper and flower petals flutter in the air, creating a sense of movement. The atmosphere is a high-detail fusion of art and reality.",
    "Seagulls soar along the seaside under the setting sun, as a couple in wedding attire holds hands.",
    "A cute, chubby werewolf holds a balloon and candy, looking adorably mischievous. The background features a full moon on a night sky.",
    "A classical beauty, dressed in a dreamy, light pink flowing gown with wide sleeves, adorned with countless tiny wind crystals.",
    "The train sped swiftly across a large bridge.",
    "In front of the door stands an apple tree with two apples glistening with dewdrops. A beautiful little bird with vibrant feathers perches on a branch, displaying intricate textures and clear details.",
]
CUSTOM_OPTION = "✍️ Enter custom prompt..."

def check_memory_usage():
    process = psutil.Process(os.getpid())
    memory_mb = process.memory_info().rss / 1024 / 1024
    print(f"🖥️ 当前内存使用: {memory_mb:.2f} MB")
    
    # 系统总内存
    total_memory = psutil.virtual_memory().total / 1024 / 1024 / 1024
    print(f"💾 系统总内存: {total_memory:.2f} GB")

def load_models():
    global pipeline, style_generator, unitok, processor, code_freq, local_repo_dir
    if "pipeline" in globals():
        return
    print('='*10, 'before download')
    local_repo_dir = snapshot_download(
        repo_id=REPO_ID,
        token=HF_TOKEN,
        allow_patterns=[
            "prior/**",          # 递归下载 prior/ 目录下所有文件
            "codebook/**",       # 递归下载 codebook/ 目录下所有文件
            "tokenizer/**",
            "processor/**",      # 递归下载 processor/ 目录下所有文件
            "text_encoder/**",   # 递归下载 text_encoder/ 目录下所有文件
            "freq.json",         # 明确指定单个文件（可选，也可用 *.json）
            "processor/**",
            "transformer/**",
            "vae/**",
            "*.json",            # 所有 .json 文件（包括 config.json 等）
            "*.pth",             # 所有 .pth 文件
            "*.safetensors",     # 所有 .safetensors 文件
        ],
        resume_download=True,
    )
    print('='*10, 'after download')
    sys.path.append(".")
    from models.pipe import CoTylePipeline
    from models.vlm_unitok import UniTok
    from models.model import StyleGenerator, Qwen2_5_VLForConditionalGeneration_Quant, Qwen2_5_VL_Quant
    from models.utils import patched_from_model_config
    from transformers import Qwen2VLProcessor, AutoConfig
    from transformers.generation.configuration_utils import GenerationConfig
    GenerationConfig.from_model_config = classmethod(patched_from_model_config)
    unitok_config = {
        "unitok_embed_dim": 3584,
        "unitok_vocab_width": 64,
        "unitok_vocab_size": 1024,
        "unitok_e_temp": 0.01,
        "unitok_num_codebooks": 1,
        "unitok_le": 0.0,
    }
    style_generator_path = os.path.join(local_repo_dir, "prior")
    config = AutoConfig.from_pretrained(style_generator_path)
    style_generator = StyleGenerator._from_config(config)
    state_dict = torch.load(os.path.join(style_generator_path, "prior.pth"), map_location="cpu")
    style_generator.load_state_dict(state_dict)
    style_generator.to(device, dtype=weight_type)
    codebook_path = os.path.join(local_repo_dir, "codebook")
    unitok = UniTok(unitok_config)
    unitok_state_dict = torch.load(os.path.join(codebook_path, "model.pth"), map_location="cpu")
    unitok.load_state_dict(unitok_state_dict)
    unitok.to(device, dtype=weight_type)
    print('='*10, 'before pipeline')
    pipeline = CoTylePipeline.from_pretrained(
        local_repo_dir,
        torch_dtype=weight_type,
        text_encoder=None,
        processor=None,
        safety_checker=None,
        requires_safety_checker=False,
        
    )

    qwen_text_visual_encoder = Qwen2_5_VLForConditionalGeneration_Quant.from_pretrained(
        local_repo_dir,
        subfolder="text_encoder",
    ).to(device, dtype=weight_type)
    qwen_text_visual_encoder = Qwen2_5_VL_Quant(unitok, qwen_text_visual_encoder)
    qwen_text_visual_encoder.to(device, dtype=weight_type)
    pipeline.text_encoder = qwen_text_visual_encoder

    processor = Qwen2VLProcessor.from_pretrained(
        local_repo_dir,
        subfolder="processor",
        min_pixels=64 * 28 * 28,
        max_pixels=256 * 28 * 28,
    )
    pipeline.processor = processor
    pipeline.to(device, dtype=weight_type)
    pipeline.set_progress_bar_config(disable=True)
    
    with open(os.path.join(local_repo_dir, "freq.json"), "r") as f:
        code_freq = json.load(f)
    print('='*10, " All models loaded successfully!")

def get_final_prompt(dropdown_val, text_val):
    if dropdown_val == CUSTOM_OPTION:
        return (text_val or "").strip()
    return (dropdown_val or "").strip()

@spaces.GPU
def generate_images(style_code, seed, num_prompts, *args):
    try:
        style_code = int(style_code)
    except Exception:
        style_code = 0
    try:
        seed = int(seed)
    except Exception:
        seed = 42
    try:
        num_prompts = int(num_prompts)
    except Exception:
        num_prompts = 1
    load_models()
    from models.utils import set_seed
    prompts = []
    for i in range(num_prompts):
        dropdown_val = args[i * 2] if i * 2 < len(args) else ""
        text_val = args[i * 2 + 1] if i * 2 + 1 < len(args) else ""
        final_prompt = get_final_prompt(dropdown_val, text_val)
        if final_prompt:
            prompts.append(final_prompt)
    if not prompts:
        raise gr.Error("Please enter at least one valid prompt!")
    set_seed(style_code)
    style_generator_inputs = {
        "input_ids": torch.randint(low=0, high=1024, size=(1, 1)).to(device),
        "attention_mask": torch.ones((1, 1)).to(device),
    }
    with torch.no_grad():
        generated_ids = style_generator.generate(
            **style_generator_inputs,
            max_new_tokens=195,
            temperature=1.0,
            top_k=200,
            top_p=0.95,
            do_sample=True,
            repetition_penalty=50.0,
            code_freq=code_freq,
            code_freq_threshold=90000,
            k=0.0001,
        )
    print('='*10, 'after style generator')
    placeholder_image = Image.new("RGB", (392, 392), (0, 0, 0))
    results = []
    for i, prompt in enumerate(prompts):
        set_seed(seed)
        inputs = {
            "image": [placeholder_image],
            "prompt": prompt,
            "generator": torch.Generator(device=device).manual_seed(seed),
            "true_cfg_scale": 6.0,
            "negative_prompt": "ugly, monster, grotesque, deformed, mutated, anatomically incorrect, distorted face, disfigured limbs, unnatural posture, blurry, low quality",
            "num_inference_steps": 40,
            "guidance_scale": 1.0,
            "num_images_per_prompt": 1,
            "codebook_id": generated_ids,
        }
        print('='*10, 'before infer')
        with torch.inference_mode():
            output = pipeline(**inputs)
        print('='*10, 'after inference')
        results.append(output.images[0])
        # output.images[0].save('tmp.png')

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()
    del output

    return results

with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
    .prompt-hint {
        font-size: 0.9em;
        color: #666;
        margin-top: -8px;
        margin-bottom: 12px;
    }
"""
) as demo:
    gr.Markdown(
        """
    <div align="center">

    ## 🎨 CoTyle: Unlocking Code-to-Style Image Generation with Discrete Style Space

    <div style="display: flex; justify-content: center; gap: 10px; flex-wrap: wrap; margin: 15px 0;">
        <a href="xxx"><img alt="Project Page" src="https://img.shields.io/badge/Project%20Page-Homepage-yellow"></a>
        <a href="xxx"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-Code-f8f0f0.svg"></a>
        <a href="xxx"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-Paper-da282a.svg"></a>
        <a href="xxxK"><img alt="Hugging Face Demo" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Demo-fd8b02"></a>
    </div>

    </div>
        """
    )

    with gr.Row():
        with gr.Column():
            style_code = gr.Number(label="Style Code", value=1234567, step=1)

            num_prompts = gr.Slider(
                minimum=1,
                maximum=6,
                value=1,
                step=1,
                label="Number of Prompts (You can choose how many prompt images to generate at once)",
            )

            all_dropdowns = []
            all_texts = []

            with gr.Column():
                for i in range(6):
                    choices = [""] + SUGGESTED_PROMPTS + [CUSTOM_OPTION]
                    dropdown = gr.Dropdown(
                        choices=choices,
                        value=SUGGESTED_PROMPTS[i] if i < len(SUGGESTED_PROMPTS) else "",
                        label=f"Prompt {i+1}",
                        interactive=True,
                        visible=(i < 1),
                    )
                    text = gr.Textbox(
                        label=f"Custom Prompt {i+1}",
                        lines=2,
                        visible=False,
                    )

                    def update_text_visibility(dropdown_val):
                        return gr.update(visible=(dropdown_val == CUSTOM_OPTION))

                    dropdown.change(
                        fn=update_text_visibility,
                        inputs=dropdown,
                        outputs=text,
                    )

                    all_dropdowns.append(dropdown)
                    all_texts.append(text)

            seed = gr.Number(label="Seed", value=42, step=1)
            run_btn = gr.Button("✨ Generate All Images", variant="primary", size="lg")

        with gr.Column():
            gallery = gr.Gallery(
                label="Generated Results",
                show_label=True,
                columns=2,
                rows=2,
                object_fit="contain",
                height="auto",
            )

    def update_components_visibility(n):
        updates = []
        for i in range(6):
            updates.append(gr.update(visible=(i < n)))
        for i in range(6):
            updates.append(gr.update(visible=False))
        return updates

    num_prompts.change(
        fn=update_components_visibility,
        inputs=num_prompts,
        outputs=(all_dropdowns + all_texts),
    )

    input_components = [style_code, seed, num_prompts]
    for d, t in zip(all_dropdowns, all_texts):
        input_components.extend([d, t])

    run_btn.click(
        fn=generate_images,
        inputs=input_components,
        outputs=gallery,
    )

    gr.Markdown(
        """
    > <strong>Tips</strong>:
    > - Adjust the <strong>Number of Prompts</strong> slider to add or remove input rows.
    > - Select <strong>"✍️ Enter custom prompt..."</strong> to type your own prompts.
    > - All images share the same `style_code`.
    """
    )


if __name__ == "__main__":
    load_models()
    demo.queue().launch(max_threads=1, share=True)