File size: 4,974 Bytes
86a2bd4
bb21a7c
 
 
 
60591c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bb21a7c
 
 
 
 
 
 
 
 
 
 
4564155
bb21a7c
 
 
 
 
 
 
e5d785e
 
 
bb21a7c
 
 
60591c0
bb21a7c
60591c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bb21a7c
60591c0
 
 
 
 
 
 
 
bb21a7c
60591c0
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
from controlnet_aux import OpenposeDetector
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
import gradio as gr
import torch
import base64
from io import BytesIO
from PIL import Image
# live conditioning
canvas_html = "<pose-canvas id='canvas-root' style='display:flex;max-width: 500px;margin: 0 auto;'></pose-canvas>"
load_js = """
async () => {
  const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/pose-gradio.js"
  fetch(url)
    .then(res => res.text())
    .then(text => {
      const script = document.createElement('script');
      script.type = "module"
      script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
      document.head.appendChild(script);
    });
}
"""
get_js_image = """
async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
  const canvasEl = document.getElementById("canvas-root");
  const data = canvasEl? canvasEl._data : null;
  return [image_in_img, prompt, image_file_live_opt, data]
}
"""

# Constants
low_threshold = 100
high_threshold = 200

# Models
pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe.enable_model_cpu_offload()

# xformers
pipe.enable_xformers_memory_efficient_attention()

# Generator seed,
generator = torch.manual_seed(0)


def get_pose(image):
    return pose_model(image)


def generate_images(image, prompt, image_file_live_opt='file', live_conditioning=None):
    if image is None and 'image' not in live_conditioning:
        raise gr.Error("Please provide an image")
    try:
        if image_file_live_opt == 'file':
            pose = get_pose(image)
        elif image_file_live_opt == 'webcam':
            base64_img = live_conditioning['image']
            image_data = base64.b64decode(base64_img.split(',')[1])
            pose = Image.open(BytesIO(image_data)).convert(
                'RGB').resize((512, 512))
        output = pipe(
            prompt,
            pose,
            generator=generator,
            num_images_per_prompt=3,
            num_inference_steps=20,
        )
        all_outputs = []
        all_outputs.append(pose)
        for image in output.images:
            all_outputs.append(image)
        return all_outputs
    except Exception as e:
        raise gr.Error(str(e))


def toggle(choice):
    if choice == "file":
        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
    elif choice == "webcam":
        return gr.update(visible=False, value=None), gr.update(visible=True, value=canvas_html)


with gr.Blocks() as blocks:
    gr.Markdown("""
        ## Generate Uncanny Faces with ControlNet Stable Diffusion
        [Check out our blog to see how this was done (and train your own controlnet)](https://huggingface.co/blog/train-your-controlnet)
    """)
    with gr.Row():
        live_conditioning = gr.JSON(value={}, visible=False)
        with gr.Column():
            image_file_live_opt = gr.Radio(["file", "webcam"], value="file",
                                           label="How would you like to upload your image?")
            image_in_img = gr.Image(source="upload", visible=True, type="pil")
            canvas = gr.HTML(None, elem_id="canvas_html", visible=False)

            image_file_live_opt.change(fn=toggle,
                                       inputs=[image_file_live_opt],
                                       outputs=[image_in_img, canvas],
                                       queue=False)
            prompt = gr.Textbox(
                label="Enter your prompt",
                max_lines=1,
                placeholder="best quality, extremely detailed",
            )
            run_button = gr.Button("Generate")
        with gr.Column():
            gallery = gr.Gallery().style(grid=[2], height="auto")
    run_button.click(fn=generate_images,
                     inputs=[image_in_img, prompt,
                             image_file_live_opt, live_conditioning],
                     outputs=[gallery],
                     _js=get_js_image)
    blocks.load(None, None, None, _js=load_js)

    gr.Examples(fn=generate_images,
                examples=[
                    ["./yoga1.jpeg",
                        "best quality, extremely detailed"]
                ],
                inputs=[image_in_img, prompt],
                outputs=[gallery],
                cache_examples=True)

blocks.launch(debug=True)