import os
from pathlib import Path
from typing import List
import torch
from PIL import Image
import gradio as gr
from ultraflux.pipeline_flux import FluxPipeline
from ultraflux.transformer_flux import FluxTransformer2DModel
from ultraflux.autoencoder_kl import AutoencoderKL

torch.set_num_threads(os.cpu_count())
torch.set_float32_matmul_precision("high")

local_vae = AutoencoderKL.from_pretrained(
    "Owen777/UltraFlux-v1",
    subfolder="vae",
    torch_dtype=torch.float32
)

transformer = FluxTransformer2DModel.from_pretrained(
    "Owen777/UltraFlux-v1-1-Transformer",
    torch_dtype=torch.float32
)

pipe = FluxPipeline.from_pretrained(
    "Owen777/UltraFlux-v1",
    vae=local_vae,
    torch_dtype=torch.float32,
    transformer=transformer
)

from diffusers import FlowMatchEulerDiscreteScheduler
pipe.scheduler.config.use_dynamic_shifting = False
pipe.scheduler.config.time_shift = 4
pipe = pipe.to("cpu")

os.makedirs("results", exist_ok=True)

def generate_ultraflux(prompt: str, seed: int = 0, steps: int = 50, size: int = 1024, guidance: float = 4.0):
    out_path = Path("results") / f"ultra_flux.png"
    with torch.inference_mode():
        image = pipe(
            prompt,
            height=size,
            width=size,
            guidance_scale=guidance,
            num_inference_steps=steps,
            max_sequence_length=512,
            generator=torch.Generator("cpu").manual_seed(seed)
        ).images[0]
    image.save(out_path)
    return out_path

demo = gr.Interface(
    fn=generate_ultraflux,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Number(label="Seed", value=0),
        gr.Slider(10, 100, step=1, value=50, label="Inference Steps"),
        gr.Slider(256, 2048, step=128, value=1024, label="Image Size"),
        gr.Slider(1.0, 10.0, step=0.1, value=4.0, label="Guidance Scale")
    ],
    outputs=gr.Image(type="filepath"),
    title="UltraFlux CPU Demo",
    description="Generate high-quality images with UltraFlux on CPU."
)

demo.launch()