import gradio as gr
import torch
from diffusers import DiffusionPipeline
import time

# Initialize the base model
base_model = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)

MAX_SEED = 2**32-1

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

def generate_image(prompt, steps, seed, cfg_scale, width, height):
    pipe.to("cuda")
    generator = torch.Generator(device="cuda").manual_seed(seed)
    
    with calculateDuration("Generating image"):
        # Generate image
        image = pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator
        ).images[0]
    return image

def run_model(prompt, cfg_scale, steps, randomize_seed, seed, width, height):
    if randomize_seed:
        seed = torch.randint(0, MAX_SEED, (1,)).item()
    
    image = generate_image(prompt, steps, seed, cfg_scale, width, height)
    return image, seed  

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Type a prompt here")
            generate_button = gr.Button("Generate")
    
    with gr.Row():
        result = gr.Image(label="Generated Image")

    with gr.Row():
        with gr.Column():
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
            randomize_seed = gr.Checkbox(True, label="Randomize seed")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)

    gr.Interface(
        fn=run_model,
        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height],
        outputs=[result, seed],
        live=True
    ).launch()