import torch
from diffusers import DiffusionPipeline
import gradio as gr

# Load the pre-trained diffusion model
pipe = DiffusionPipeline.from_pretrained('ptx0/terminus-xl-velocity-v2', torch_dtype=torch.bfloat16)
pipe.to('cuda')

# Define the image generation function with adjustable parameters and a progress bar
def generate(prompt, guidance_scale, num_inference_steps, negative_prompt):
    with gr.Progress(steps=num_inference_steps) as progress:
        for i in range(num_inference_steps):
            progress.update(progress=i)
        return pipe(
            prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps
        ).images

# Example prompts to demonstrate the model's capabilities
example_prompts = [
    ["A futuristic cityscape at night under a starry sky", 7.5, 25, "blurry, overexposed"],
    ["A serene landscape with a flowing river and autumn trees", 8.0, 20, "crowded, noisy"],
    ["An abstract painting of joy and energy in bright colors", 9.0, 30, "dark, dull"]
]

# Create a Gradio interface
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Enter your prompt"),
        gr.Slider(5, 10, step=0.1, label="Guidance Scale", value=7.5),
        gr.Slider(10, 50, step=5, label="Number of Inference Steps", value=25),
        gr.Text(value="underexposed, blurry, ugly, washed-out", label="Negative Prompt")
    ],
    outputs=gr.Gallery(height=1024, min_width=1024, columns=2),
    examples=example_prompts,
    title="Image Generation with Diffusion Model",
    description="Generate images based on textual prompts. Adjust the parameters to see how they affect the outcome."
).launch()