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

# Load the pre-trained diffusion model
pipe = StableDiffusion3Pipeline.from_pretrained('ptx0/sd3-reality-mix', torch_dtype=torch.bfloat16)
pipe.to('cuda')
import re

def extract_resolution(resolution_str):
    match = re.match(r'(\d+)x(\d+)', resolution_str)
    if match:
        width = int(match.group(1))
        height = int(match.group(2))
        return (width, height)
    else:
        return None

# Define the image generation function with adjustable parameters and a progress bar
@spaces.GPU
def generate(prompt, guidance_scale, num_inference_steps, resolution, negative_prompt):
    width, height = extract_resolution(resolution) or (1024, 1024)
    return pipe(
        prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width, height=height
    ).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, 1024x1024,1152x960,896x1152
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Enter your prompt"),
        gr.Slider(1, 20, step=0.1, label="Guidance Scale", value=3.4),
        gr.Slider(1, 50, step=1, label="Number of Inference Steps", value=28),
        gr.Radio(["1024x1024", "1152x960", "896x1152"], label="Resolution", value="1152x960"),
        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="SD3 Reality Mix",
    description="SD3 Reality Mix is an attempt to finetune SD3 Medium 2B on photo-concept-bucket, a 576k photo set of images."
).launch()