print('v4')
import os
import torch
from PIL import Image
from io import BytesIO
import json
from huggingface_hub import login, hf_hub_download
import spaces
import gradio as gr
token=os.environ.get("HF_TOKEN")

login(token=os.environ.get("HF_TOKEN"))
REPO_ID = "Kwai-Kolors/cotyle"
# Use GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
weight_type = torch.bfloat16 if device == "cuda" else torch.float32

# Predefined suggested prompts (already in English)
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..."

# Lazy load models to avoid slow startup
def load_models():
    global pipeline, style_generator, unitok, processor, code_freq
    if 'pipeline' in globals():
        return  # Already loaded

    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 set_seed, patched_from_model_config
    from transformers import Qwen2VLProcessor
    from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
    from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
    from transformers.generation.configuration_utils import GenerationConfig
    _original_from_model_config = GenerationConfig.from_model_config
    GenerationConfig.from_model_config = classmethod(patched_from_model_config)

    model_path = "Kwai-Kolors/cotyle"
    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
    }

    # Load Style Generator
    style_generator_path = hf_hub_download(
        repo_id=model_path,
        filename='prior',
        token=token,
    )

    from transformers import AutoConfig
    config = AutoConfig.from_pretrained(f"{style_generator_path}/config.json")
    style_generator = StyleGenerator._from_config(config)
    state_dict = torch.load(f"{style_generator_path}/prior.pth", map_location='cpu')
    style_generator.load_state_dict(state_dict)
    style_generator.to(device, dtype=weight_type)

    # Load UniTok

    codebook_path = hf_hub_download(
        repo_id=model_path,
        filename='codebook',
        token=token,
    )

    unitok = UniTok(unitok_config)
    unitok_state_dict = torch.load(f"{codebook_path}/model.pth", map_location='cpu')
    unitok.load_state_dict(unitok_state_dict)
    unitok.to(device, dtype=weight_type)

    # Load Pipeline (without text encoder initially)
    pipeline = CoTylePipeline.from_pretrained(
        model_path,
        torch_dtype=weight_type,
        text_encoder=None,
        processor=None,
        safety_checker=None,
        requires_safety_checker=False
    )

    # Load Qwen2.5-VL Text-Visual Encoder
    from transformers import Qwen2_5_VLForConditionalGeneration
    qwen_text_visual_encoder = Qwen2_5_VLForConditionalGeneration_Quant.from_pretrained(
    model_path,
    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

    # Load Processor
    processor = Qwen2VLProcessor.from_pretrained(
        model_path, 
        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)

    # Load code frequency
    with open(f'{model_path}/freq.json', 'r') as f:
        code_freq = json.load(f)

    print("✅ All models loaded successfully!")


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

@spaces.GPU
def generate_images(style_code: int, seed: int, num_prompts: int, *args):
    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!")

    # Step 1: Generate style codebook tokens
    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,
        )

    # Step 2: Generate images
    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,
        }

        with torch.inference_mode():
            output = pipeline(**inputs)
        results.append(output.images[0])

    return results


# Gradio Interface
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

    Enter a `style code` and multiple prompts to generate stylized images.

    <p align="center">
    <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>
    </p>

    </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=4, 
                step=1, 
                label="Number of Prompts (You can choose how many prompt images to generate at once)"
            )

            all_dropdowns = []
            all_texts = []
            prompt_rows = []
            
            with gr.Column():
                for i in range(6):
                    with gr.Row(visible=(i < 4)) as row:
                        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
                        )
                        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)
                        prompt_rows.append(row)
            
            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"
            )

    # Update visibility of prompt rows
    def update_rows_visibility(n):
        return [gr.update(visible=(i < n)) for i in range(6)]
    
    num_prompts.change(
        fn=update_rows_visibility,
        inputs=num_prompts,
        outputs=prompt_rows
    )

    # Build input list: [style_code, seed, num_prompts, d1, t1, d2, t2, ...]
    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("""
    > **Tips**:
    > - Adjust the **Number of Prompts** slider to add or remove input rows.
    > - Select **"✍️ Enter custom prompt..."** to type your own prompts.
    > - All images share the same `style_code`.
    """)

# Launch
if __name__ == "__main__":
    import sys
    sys.path.append(".")
    demo.queue.launch(debug=True)