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

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

MODEL_ID = "drbaph/FLUX.1-schnell-dev-merged-fp8-4step"
MODEL_FILE = "flux1-schnell-dev-merged-fp8-4step.safetensors"

def load_model():
    pipe = FluxPipeline.from_single_file(
        f"https://huggingface.co/{MODEL_ID}/resolve/main/{MODEL_FILE}",
        torch_dtype=dtype
    )
    pipe.to(device)
    return pipe

pipe = load_model()

MAX_SEED = 2**32 - 1
MAX_IMAGE_SIZE = 2048

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = torch.randint(0, MAX_SEED, (1,)).item()
    generator = torch.Generator(device=device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0,
        max_sequence_length=256
    ).images[0]
    return image, seed

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# FLUX.1 [schnell-dev-merged-fp8-4step]")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt")
        run_button = gr.Button("Generate")
    with gr.Row():
        result = gr.Image(label="Generated Image")
        seed_output = gr.Number(label="Seed Used")
    with gr.Accordion("Advanced Settings", open=False):
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
        num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=10, step=1, value=4)
    
    inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps]
    run_button.click(fn=infer, inputs=inputs, outputs=[result, seed_output])

demo.launch()