"""
"""
from datetime import datetime

# Upgrade PyTorch
import os
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 torch torchvision spaces')

# CUDA toolkit install
t0 = datetime.now()
from utils.cuda_toolkit import install_cuda_toolkit; install_cuda_toolkit()
print('install_cuda_toolkit', -(t0 - (t0 := datetime.now())))

# Actual app.py
import os

import gradio as gr
import spaces
import torch
import torch._inductor
from diffusers import FluxPipeline

from .utils.zerogpu import aoti_compile


pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
print('FluxPipeline.from_pretrained', -(t0 - (t0 := datetime.now())))
package_path = 'pipeline.pt2'


@spaces.GPU(duration=1500)
def compile_transformer():

    def _example_tensor(*shape):
        return torch.randn(*shape, device='cuda', dtype=torch.bfloat16)

    is_timestep_distilled = not pipeline.transformer.config.guidance_embeds
    seq_length = 256 if is_timestep_distilled else 512

    transformer_kwargs = {
        'hidden_states': _example_tensor(1, 4096, 64),
        'timestep': torch.tensor([1.], device='cuda', dtype=torch.bfloat16),
        'guidance': None if is_timestep_distilled else torch.tensor([1.], device='cuda', dtype=torch.bfloat16),
        'pooled_projections': _example_tensor(1, 768),
        'encoder_hidden_states': _example_tensor(1, seq_length, 4096),
        'txt_ids': _example_tensor(seq_length, 3),
        'img_ids': _example_tensor(4096, 3),
        'joint_attention_kwargs': {},
        'return_dict': False,
    }

    inductor_configs = {
        'conv_1x1_as_mm': True,
        'epilogue_fusion': False,
        'coordinate_descent_tuning': True,
        'coordinate_descent_check_all_directions': True,
        'max_autotune': True,
        'triton.cudagraphs': True,
    }

    exported = torch.export.export(pipeline.transformer, args=(), kwargs=transformer_kwargs)

    return aoti_compile(exported, inductor_configs)


transformer_config = pipeline.transformer.config
pipeline.transformer = compile_transformer()
pipeline.transformer.config = transformer_config

@spaces.GPU
def _generate_image(prompt: str, t0: datetime):
    print('@spaces.GPU', -(t0 - (t0 := datetime.now())))
    images = []
    for _ in range(4):
        images += pipeline(prompt, num_inference_steps=4).images
        print('pipeline', -(t0 - (t0 := datetime.now())))
    return images


def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
    return _generate_image(prompt, datetime.now())

gr.Interface(generate_image, gr.Text(), gr.Gallery()).launch(show_error=True)