app.py

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


seed = None
if seed is None:
    seed = int.from_bytes(os.urandom(2), "big")
print(f"Using seed: {seed}")
generator = torch.Generator(device=device).manual_seed(seed)

def generate_image(prompt):
    with torch.device(device):
        print(f"Current device: {torch.device(device)}")
        image = pipe(
            prompt=prompt,
            height=(height := 512),
            width=(width := 512),
            num_inference_steps=28,
            guidance_scale=7.0,
            num_images_per_prompt=1,
            generator=generator,
            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()