import gradio as gr
from diffusers import StableDiffusion3Pipeline
import torch
import os
import psutil

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.float16,
)

def generate_image(prompt):
    with torch.no_grad():
        image = pipe(
            prompt=prompt,
            height=(height := 512),
            width=(width := 512),
            num_inference_steps=28,
            guidance_scale=7.0,
            num_images_per_prompt=1,
            output_type="pil",
            return_dict=True,
            callback_on_step_end_tensor_inputs=["latents"],
        ).images[0]
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text", 
    outputs="image", 
    title="Stable Diffusion 3.5",
    description="Enter a prompt to generate an image using Stable Diffusion 3.5."
)

iface.launch()