File size: 3,473 Bytes
fe01e72 dd0f699 fe01e72 a9c4647 fe01e72 47500b9 fe01e72 052501f fe01e72 a9c4647 fe01e72 a9c4647 fe01e72 a9c4647 fe01e72 5a39d55 052501f fe01e72 c68fe52 fe01e72 052501f c68fe52 fe01e72 a9c4647 fe01e72 c68fe52 fe01e72 | 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 | import gradio as gr
import spaces
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import numpy as np
import imageio
import tempfile
import os
# stabilityai/stable-diffusion-2 was withdrawn from the Hub.
# This is the community mirror of the same checkpoint.
MODEL_ID = "sd2-community/stable-diffusion-2"
# Global pipeline variable
pipe = None
def initialize_pipeline():
"""Initialize the pipeline if not already loaded."""
global pipe
if pipe is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Initializing pipeline on device: {device}")
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
)
pipe = pipe.to(device)
return pipe
@spaces.GPU
def generate_image(prompt, seed, num_inference_steps):
# Initialize pipeline
pipeline = initialize_pipeline()
device = pipeline.device
# Set the random seed for reproducibility
generator = torch.Generator(device=device).manual_seed(int(seed))
# Store intermediate images
frames = []
def callback(step: int, timestep: int, latents):
# Decode latents to image
with torch.no_grad():
image = pipeline.decode_latents(latents)
image = pipeline.numpy_to_pil(image)[0]
frames.append(image)
# Generate the image with callback for each step
with torch.no_grad():
result = pipeline(
prompt=prompt,
num_inference_steps=int(num_inference_steps),
generator=generator,
callback=callback,
callback_steps=1,
)
# Save frames as video
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
video_path = tmpfile.name
imageio.mimsave(video_path, frames, fps=5)
# Return final image and video path
return result.images[0], video_path
def create_interface():
"""Create and configure the Gradio interface."""
# Create the Gradio interface
interface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(
label="Prompt",
placeholder="Enter a text description of the image you want to generate...",
lines=3,
),
gr.Slider(
minimum=0,
maximum=1000000,
randomize=True,
step=1,
label="Seed",
info="Random seed. Keep the same seed to generate the same image or leave it random to generate a new image each time.",
),
gr.Slider(
minimum=1,
maximum=50,
value=30,
step=1,
label="Diffusion Steps",
info="Number of denoising steps (more steps = higher quality but slower)",
),
],
outputs=[
gr.Image(label="Generated Image", type="pil"),
gr.Video(label="Diffusion Steps Video"),
],
title="Stable Diffusion Image Generator",
description="Generates an image using Stable Diffusion 2 showing the diffusion process as a video.",
)
return interface
if __name__ == "__main__":
# Create and launch the interface
demo = create_interface()
demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
|