Spaces:
Runtime error
Runtime error
import torch | |
from diffusers import StableDiffusion3Pipeline | |
import gradio as gr | |
import spaces | |
# 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 | |
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() | |