Spaces:
Running
Running
import torch | |
from diffusers import StableDiffusionPipeline | |
from controlnet_aux import CannyDetector, OpenposeDetector, MidasDetector | |
import gradio as gr | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16) | |
if torch.backends.mps.is_available(): | |
device = "mps" | |
elif torch.cuda.is_available(): | |
device = "cuda" | |
else: | |
device = "cpu" | |
pipe.to(device) | |
canny_detector = CannyDetector() | |
pose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet") | |
midas_detector = MidasDetector.from_pretrained("lllyasviel/ControlNet") | |
def generate_pics(prompt): | |
image = pipe(prompt).images[0] | |
canny_image = canny_detector(image) | |
pose_image = pose_detector(image) | |
depth_image = midas_detector(image) | |
return image, canny_image, pose_image, depth_image | |
gr.Interface( | |
fn=generate_pics, | |
inputs=gr.Textbox(lines=2, label="Enter your prompt"), | |
outputs=[ | |
gr.Image(label="Generated Image"), | |
gr.Image(label="Canny"), | |
gr.Image(label="OpenPose"), | |
gr.Image(label="Depth"), | |
], | |
title="Image Generation using Stable Diffusion", | |
description="Enter the prompt to generate an image using Stable Diffusion" | |
).launch() | |