import math
import os
import random

import gradio as gr
import numpy as np
import spaces
import torch

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from diffusers import Flux2KleinPipeline
from lakonlab.models.architectures import OklabColorEncoder
from lakonlab.models.diffusions.schedulers import FlowAdapterScheduler
from lakonlab.pipelines.pipeline_pixelflux2_klein import PixelFlux2KleinPipeline
from lakonlab.ui.gradio.create_text_to_img import create_interface_text_to_img

from huggingface_hub import login
login(token=os.getenv('HF_TOKEN'))


DEFAULT_PROMPT = (
    'Restored color photo from the 1900s. A middle-aged man with cybernetic metal hands is sitting on an old wooden '
    'chair and reading the newspaper. The newspaper has the prominent headline "AsymFLOW RELEASED" in large bold font. '
    'Close-up shot focusing on the newspaper.'
)
DEFAULT_NEG_PROMPT = 'Low quality, worst quality, blurry, deformed, bad anatomy, unclear text'


def set_random_seed(seed: int, deterministic: bool = True) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


base_pipe = Flux2KleinPipeline.from_pretrained(
    'black-forest-labs/FLUX.2-klein-base-9B',
    vae=None,
    scheduler=None,
    torch_dtype=torch.bfloat16)
base_pipe = base_pipe.to('cuda')
scheduler = FlowAdapterScheduler(
    shift=17.0,
    use_dynamic_shifting=True,
    base_seq_len=1024 ** 2,
    max_seq_len=2048 ** 2,
    base_logshift=math.log(17.0),
    max_logshift=math.log(34.0),
    dynamic_shifting_type='sqrt',
    base_scheduler='UniPCMultistep')
color_encoder = OklabColorEncoder(
    use_affine_norm=True,
    mean=(0.56, 0.0, 0.01),
    std=0.16).to('cuda')

pipe = PixelFlux2KleinPipeline(
    scheduler=scheduler,
    vae=color_encoder,
    text_encoder=base_pipe.text_encoder,
    tokenizer=base_pipe.tokenizer,
    transformer=base_pipe.transformer)
pipe.load_lakonlab_adapter(
    'Lakonik/AsymFLUX.2-klein-9B',
    target_module_name='transformer')

pipe_sft_zimage_turbo = PixelFlux2KleinPipeline(
    scheduler=scheduler,
    vae=color_encoder,
    text_encoder=base_pipe.text_encoder,
    tokenizer=base_pipe.tokenizer,
    transformer=base_pipe.transformer)
pipe_sft_zimage_turbo.load_lakonlab_adapter(
    'Lakonik/AsymFLUX.2-klein-9B-collection',
    subfolder='asymflux2_klein_9b_sft_zimage_turbo',
    target_module_name='transformer')

pipe_sft_flux2_klein = PixelFlux2KleinPipeline(
    scheduler=scheduler,
    vae=color_encoder,
    text_encoder=base_pipe.text_encoder,
    tokenizer=base_pipe.tokenizer,
    transformer=base_pipe.transformer)
pipe_sft_flux2_klein.load_lakonlab_adapter(
    'Lakonik/AsymFLUX.2-klein-9B-collection',
    subfolder='asymflux2_klein_9b_sft_flux2_klein',
    target_module_name='transformer')

del base_pipe

model_pipes = {
    'AsymFLUX.2 klein 9B (base)': pipe,
    'AsymFLUX.2 klein 9B SFT Z-Image Turbo (finetuned on synthetic data generated by Z-Image Turbo)': pipe_sft_zimage_turbo,
    'AsymFLUX.2 klein 9B SFT FLUX.2 klein (finetuned on synthetic data generated by FLUX.2 klein Distilled 9B)': pipe_sft_flux2_klein,
}
default_model = list(model_pipes.keys())[1]


@spaces.GPU
def generate(
        model_name, seed, prompt, negative_prompt, width, height, steps, guidance_scale,
        progress=gr.Progress(track_tqdm=True)):
    selected_pipe = model_pipes[model_name]
    return selected_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        generator=torch.Generator().manual_seed(seed),
    ).images[0]


with gr.Blocks(
        analytics_enabled=False,
        title='AsymFLUX.2 klein Demo',
        css_paths='lakonlab/ui/gradio/style.css') as demo:
    gr.Markdown(
        '# AsymFLUX.2 klein Demo\n\n'
        'Pixel-space text-to-image generation demo of the paper '
        '[Asymmetric Flow Models](https://arxiv.org/abs/2605.12964). '
        '**Base model:** [FLUX.2 klein Base 9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B). '
        '**Code:** [https://github.com/Lakonik/LakonLab](https://github.com/Lakonik/LakonLab).\n'
        '<br> Use and distribution of this app are governed by the '
        '[FLUX Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B/blob/main/LICENSE.md).'
    )

    model_dropdown = gr.Dropdown(
        choices=list(model_pipes.keys()),
        value=default_model,
        label='Select model variant',
        elem_classes=['force-hide-container'])

    create_interface_text_to_img(
        generate,
        prompt=DEFAULT_PROMPT,
        negative_prompt=DEFAULT_NEG_PROMPT,
        steps=38,
        min_steps=4,
        max_steps=50,
        guidance_scale=4.0,
        height=1280,
        width=960,
        create_negative_prompt=True,
        args=[
            model_dropdown,
            'last_seed',
            'prompt',
            'negative_prompt',
            'width',
            'height',
            'steps',
            'guidance_scale',
        ])

demo.queue().launch()