File size: 5,163 Bytes
7a5c6a0
 
 
 
 
 
 
 
 
 
 
8b2a350
7a5c6a0
 
 
efe3c52
7a5c6a0
efe3c52
 
7a5c6a0
efe3c52
 
 
2e755f1
 
3b1f734
efe3c52
 
3b1f734
 
 
 
 
7a5c6a0
c5a40b1
 
3b1f734
 
c5a40b1
7a5c6a0
 
 
 
 
 
 
97f34a9
7a5c6a0
 
 
 
 
efe3c52
 
7a5c6a0
 
71b8b82
7a5c6a0
 
 
efe3c52
7a5c6a0
 
 
 
 
8b2a350
 
7a5c6a0
c5a40b1
3b1f734
c5a40b1
7a5c6a0
c5a40b1
7a5c6a0
 
 
ad64d06
7a5c6a0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c5a40b1
7a5c6a0
 
c5a40b1
d44dcbe
7a5c6a0
 
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
import cv2
import einops
import gradio as gr
import numpy as np
import torch


from pytorch_lightning import seed_everything
from util import resize_image, HWC3, apply_canny
from ldm.models.diffusion.ddim import DDIMSampler


from cldm.model import create_model, load_state_dict

from huggingface_hub import hf_hub_url, cached_download

REPO_ID = "lllyasviel/ControlNet"
canny_checkpoint = "models/control_sd15_canny.pth"
pose_checkpoint = "models/control_sd15_openpose.pth"

canny_model = create_model('./models/cldm_v15.yaml')
canny_model.load_state_dict(load_state_dict(cached_download(
    hf_hub_url(REPO_ID, canny_checkpoint)
), location='cpu'))
canny_model = canny_model.cuda()
ddim_sampler = DDIMSampler(canny_model)


# pose_model = create_model('./models/cldm_v15.yaml')
# pose_model.load_state_dict(load_state_dict(cached_download(
#     hf_hub_url(REPO_ID, pose_checkpoint)
# ), location='cpu'))
# ddim_sampler_pose = DDIMSampler(pose_model)

def process(input_image, prompt, input_control, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
    # TODO: Add other control tasks
    return process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold)
    
def process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
    with torch.no_grad():
        img = resize_image(HWC3(input_image), image_resolution)
        H, W, C = img.shape

        detected_map = apply_canny(img, low_threshold, high_threshold)
        detected_map = HWC3(detected_map)

        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
        control = torch.stack([control for _ in range(num_samples)], dim=0)
        control = einops.rearrange(control, 'b h w c -> b c h w').clone()

        seed_everything(seed)

        cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
        un_cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([n_prompt] * num_samples)]}
        shape = (4, H // 8, W // 8)

        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
                                                     shape, cond, verbose=False, eta=eta,
                                                     unconditional_guidance_scale=scale,
                                                     unconditional_conditioning=un_cond)
        x_samples = canny_model.decode_first_stage(samples)
        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)

        results = [x_samples[i] for i in range(num_samples)]
    return [255 - detected_map] + results


    
block = gr.Blocks().queue()
control_task_list = [
    "Canny Edge Map"
]
with block:
    gr.Markdown("## Adding Conditional Control to Text-to-Image Diffusion Models")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(source='upload', type="numpy")
            input_control = gr.Dropdown(control_task_list, value="Canny Edge Map", label="Control Task")
            prompt = gr.Textbox(label="Prompt")
            run_button = gr.Button(label="Run")
            with gr.Accordion("Advanced options", open=False):
                num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=256)
                low_threshold = gr.Slider(label="Canny low threshold", minimum=1, maximum=255, value=100, step=1)
                high_threshold = gr.Slider(label="Canny high threshold", minimum=1, maximum=255, value=200, step=1)
                ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
                scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
                seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
                eta = gr.Number(label="eta (DDIM)", value=0.0)
                a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
                n_prompt = gr.Textbox(label="Negative Prompt",
                                      value='longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair,extra digit, fewer digits, cropped, worst quality, low quality')
        with gr.Column():
            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
    ips = [input_image, prompt, input_control, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold]
    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])

    examples = gr.Examples(examples=[["bird.png", "bird","Canny Edge Map"]],inputs = [input_image, prompt, input_control], outputs = [result_gallery])


block.launch(debug = True)