Spaces:
Runtime error
Runtime error
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) |