Spaces:
Running
Running
File size: 5,061 Bytes
5b7dd46 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
import os
import random
import time
from typing import Optional
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
# -----------------------------
# Device & Precision
# -----------------------------
USE_CUDA = torch.cuda.is_available()
DTYPE = torch.float16 if USE_CUDA else torch.float32
DEVICE = "cuda" if USE_CUDA else "cpu"
MODEL_ID = os.environ.get("MODEL_ID", "runwayml/stable-diffusion-v1-5")
pipe: Optional[StableDiffusionPipeline] = None
def load_pipeline():
"""Load and configure the Stable Diffusion pipeline once at startup."""
global pipe
t0 = time.time()
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_ID,
torch_dtype=DTYPE,
safety_checker=None, # Keep None for faster demos
)
# Use a fast, good-quality scheduler
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(DEVICE)
# Optional memory optimization on GPU
if USE_CUDA:
try:
pipe.enable_attention_slicing()
pipe.enable_xformers_memory_efficient_attention()
except Exception:
pass
t1 = time.time()
print(f"Pipeline loaded in {t1 - t0:.2f}s on {DEVICE} (dtype={DTYPE}).")
# Load on import (Space boot)
load_pipeline()
def generate_image(
prompt: str,
negative_prompt: str,
steps: int,
guidance: float,
width: int,
height: int,
seed: int,
):
if not prompt or len(prompt.strip()) == 0:
raise gr.Error("Please enter a prompt.")
width = max(256, min(1024, width))
height = max(256, min(1024, height))
if seed == -1:
seed = random.randint(0, 2**31 - 1)
generator = torch.Generator(device=DEVICE).manual_seed(seed)
with torch.autocast(DEVICE, enabled=USE_CUDA):
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt or None,
num_inference_steps=int(steps),
guidance_scale=float(guidance),
width=int(width),
height=int(height),
generator=generator,
).images[0]
return image, seed
# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks(title="Stable Diffusion Image Generator", css="footer {visibility: hidden}") as demo:
gr.Markdown(
"""
# 🧠 Stable Diffusion Image Generator
Type a prompt and generate an image using **Stable Diffusion v1.5**.
**Tip:** For consistent results, set a fixed seed. Use `-1` for random seed.
"""
)
with gr.Row():
with gr.Column(scale=3):
prompt = gr.Textbox(
label="Prompt",
placeholder="a cinematic portrait of an astronaut relaxing in a tropical cafe, 35mm photo, bokeh, soft light",
lines=3,
)
negative_prompt = gr.Textbox(
label="Negative Prompt (optional)",
placeholder="blurry, low quality, extra fingers, text, watermark",
lines=2,
)
with gr.Row():
steps = gr.Slider(5, 50, value=25, step=1, label="Steps")
guidance = gr.Slider(0.0, 15.0, value=7.5, step=0.5, label="Guidance Scale")
with gr.Row():
width = gr.Slider(256, 1024, value=512, step=64, label="Width")
height = gr.Slider(256, 1024, value=512, step=64, label="Height")
seed = gr.Number(value=-1, precision=0, label="Seed (-1 for random)")
generate_btn = gr.Button("Generate", variant="primary")
with gr.Column(scale=4):
out_image = gr.Image(label="Result", type="pil")
out_seed = gr.Number(label="Used Seed", interactive=False)
examples = gr.Examples(
examples=[
[
"ultra-detailed watercolor of a koi fish swirling through clouds, ethereal, pastel palette",
"lowres, noisy, text",
28,
7.5,
512,
512,
1234,
],
[
"cozy cyberpunk alley coffee shop at dusk, volumetric lighting, rain reflections, 4k",
"low quality, oversaturated",
25,
6.5,
640,
384,
-1,
],
[
"studio photo of a cute corgi wearing sunglasses, soft light, shallow depth of field",
"text, watermark, blurry",
22,
7.0,
512,
512,
2024,
],
],
inputs=[prompt, negative_prompt, steps, guidance, width, height, seed],
)
generate_btn.click(
fn=generate_image,
inputs=[prompt, negative_prompt, steps, guidance, width, height, seed],
outputs=[out_image, out_seed],
api_name="generate",
)
if __name__ == "__main__":
demo.launch()
|