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


# 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)

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")


@spaces.GPU(duration=70)
def generate_image(
    prompt_mash,
    width,
    height,
):
    pipe.to("cuda")
    if negative_prompt == '':
        negative_prompt =  "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"


    seed = 235234
    num_images = 1
    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=4,
                true_cfg_scale=1,
                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, 
    width, 
    height, 
    progress=gr.Progress(track_tqdm=True)
):

    with calculateDuration("Loading Lightning LoRA and style LoRA"):
        pipe.load_lora_weights(
            'Wuli-Art/Qwen-Image-2512-Turbo-LoRA', 
            weight_name='Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V1.0-bf16.safetensors',
            adapter_name="lightning"
        )
        pipe.set_adapters(["lightning"], adapter_weights=[1.0])

    
    multiplier = float(quality_multiplier.replace('x', ''))
    width = int(width * multiplier)
    height = int(height * multiplier)
    num_images = int(quantity) + 1


    pairs = generate_image(
        prompt,
        width,
        height
    )

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

    return images_for_gallery



css = '''
#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}
#speed_status{padding: .5em; border-radius: 5px; margin: 1em 0}
'''


with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 60)) as app:
    title = gr.HTML(
        """<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\">        
        <h3 style=\"margin-top: -10px\">Wuli-art/Qwen-Image-2512-Turbo-LoRA</h3>""",
        elem_id="title",
    )
    
    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")
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn", interactive=False)
        with gr.Column(scale=1, elem_id="gen_column"):
            result = gr.Gallery(label="Generated Images", show_label=True, elem_id="result_gallery")
    

    generate_event = gr.on(
        triggers=[generate_button.click, prompt.submit],
        fn=run_lora,
        inputs=[
            prompt, 
            width, 
            height
        ],
        outputs=[result]
    )


app.queue()
app.launch()