File size: 1,273 Bytes
6fa2c62
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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()