Spaces:

kratadata
/

diffusion-steps-visualiser

Sleeping

Paulina

fix endpoint

c68fe52 13 days ago

3.47 kB

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)

	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)