app.py

import os
import sys
sys.path.append(os.getcwd())
sys.path.append(os.path.join(os.getcwd(), "annotator/entityseg"))
import cv2
import spaces
import einops
import torch
import gradio as gr
import numpy as np
from pytorch_lightning import seed_everything
from PIL import Image

from annotator.util import resize_image, HWC3
from annotator.canny import CannyDetector
from annotator.midas import MidasDetector
from annotator.entityseg import EntitysegDetector
from annotator.openpose import OpenposeDetector
from annotator.content import ContentDetector
from annotator.cielab import CIELabDetector

from models.util import create_model, load_state_dict
from models.ddim_hacked import DDIMSampler

'''
define conditions
'''
max_conditions = 8
condition_types = ["edge", "depth", "seg", "pose", "content", "color"]

apply_canny = CannyDetector()
apply_midas = MidasDetector()
apply_seg = EntitysegDetector()
apply_openpose = OpenposeDetector()
apply_content = ContentDetector()
apply_color = CIELabDetector()

processors = {
 "edge": apply_canny,
 "depth": apply_midas,
 "seg": apply_seg,
 "pose": apply_openpose,
 "content": apply_content,
 "color": apply_color,
}

descriptors = {
 "edge": "canny",
 "depth": "depth",
 "seg": "segmentation",
 "pose": "openpose",
}


@torch.no_grad()
def get_unconditional_global(c_global):
 if isinstance(c_global, dict):
 return {k:torch.zeros_like(v) for k,v in c_global.items()}
 elif isinstance(c_global, list):
 return [torch.zeros_like(c) for c in c_global]
 else:
 return torch.zeros_like(c_global) 


@spaces.GPU
def process(prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, 
 strength, scale, seed, eta, global_strength, color_strength, local_strength, *args):
 
 seed_everything(seed)

 conds_and_types = args
 conds = conds_and_types[0::2]
 types = conds_and_types[1::2]
 conds = [c for c in conds if c is not None]
 types = [t for t in types if t is not None]
 assert len(conds) == len(types)

 detected_maps = []
 other_maps = []
 tasks = []

 # initialize global control
 global_conditions = dict(clipembedding=np.zeros((1, 768), dtype=np.float32), color=np.zeros((1, 180), dtype=np.float32))
 global_control = {}
 for key in global_conditions.keys():
 global_cond = torch.from_numpy(global_conditions[key]).unsqueeze(0).repeat(num_samples, 1, 1)
 global_cond = global_cond.cuda().to(memory_format=torch.contiguous_format).float()
 global_control[key] = global_cond

 # initialize local control
 anchor_image = HWC3(np.zeros((image_resolution, image_resolution, 3)).astype(np.uint8))
 oH, oW = anchor_image.shape[:2]
 H, W, C = resize_image(anchor_image, image_resolution).shape
 anchor_tensor = ddim_sampler.model.qformer_vis_processor['eval'](Image.fromarray(anchor_image))
 local_control = torch.tensor(anchor_tensor).cuda().to(memory_format=torch.contiguous_format).half()

 task_prompt = ''

 with torch.no_grad():

 # set up local control
 for cond, typ in zip(conds, types):
 if typ in ['edge', 'depth', 'seg', 'pose']:
 oH, oW = cond.shape[:2] 
 cond_image = HWC3(cv2.resize(cond, (W, H)))
 cond_detected_map = processors[typ](cond_image)
 cond_detected_map = HWC3(cond_detected_map)
 detected_maps.append(cond_detected_map)
 tasks.append(descriptors[typ])
 elif typ in ['content']:
 other_maps.append(cond)
 content_image = cv2.cvtColor(cond, cv2.COLOR_RGB2BGR)
 content_emb = apply_content(content_image)
 global_conditions['clipembedding'] = content_emb
 elif typ in ['color']:
 color_hist = apply_color(cond)
 global_conditions['color'] = color_hist
 color_palette = apply_color.hist_to_palette(color_hist) # (50, 189, 3)
 color_palette = cv2.resize(color_palette, (W, H), cv2.INTER_NEAREST)
 other_maps.append(color_palette)
 if len(detected_maps) > 0:
 local_control = torch.cat([ddim_sampler.model.qformer_vis_processor['eval'](Image.fromarray(img)).cuda().unsqueeze(0) for img in detected_maps], dim=1)
 task_prompt = ' conditioned on ' + ' and '.join(tasks)
 local_control = local_control.repeat(num_samples, 1, 1, 1)

 # set up global control
 for key in global_conditions.keys():
 global_cond = torch.from_numpy(global_conditions[key]).unsqueeze(0).repeat(num_samples, 1, 1)
 global_cond = global_cond.cuda().to(memory_format=torch.contiguous_format).float()
 global_control[key] = global_cond

 # set up prompt
 input_prompt = (prompt + ' ' + task_prompt).strip()

 # set up cfg
 uc_local_control = local_control
 uc_global_control = get_unconditional_global(global_control)
 cond = {
 "local_control": [local_control], 
 "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)], 
 "global_control": [global_control],
 "text": [[input_prompt] * num_samples],
 }
 un_cond = {
 "local_control": [uc_local_control], 
 "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)], 
 'global_control': [uc_global_control],
 "text": [[input_prompt] * num_samples],
 }
 shape = (4, H // 8, W // 8)

 model.control_scales = [strength] * 13
 samples, _ = ddim_sampler.sample(ddim_steps, num_samples,
 shape, cond, verbose=False, eta=eta,
 unconditional_guidance_scale=scale,
 unconditional_conditioning=un_cond, 
 global_strength=global_strength, 
 color_strength=color_strength,
 local_strength=local_strength)

 x_samples = 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)]

 results = [cv2.resize(res, (oW, oH)) for res in results]
 detected_maps = [cv2.resize(maps, (oW, oH)) for maps in detected_maps]
 return [results, detected_maps+other_maps]

 
def variable_image_outputs(k):
 if k is None:
 k = 1
 k = int(k)
 imageboxes = []
 for i in range(max_conditions):
 if i<k:
 with gr.Row(visible=True):
 img = gr.Image(sources=['upload'], type="numpy", label=f'Condition {i+1}', visible=True, interactive=True, scale=3, height=200)
 typ = gr.Dropdown(condition_types, visible=True, interactive=True, label="type", scale=1)
 else:
 with gr.Row(visible=False):
 img = gr.Image(sources=['upload'], type="numpy", label=f'Condition {i+1}', visible=False, scale=3, height=200)
 typ = gr.Dropdown(condition_types, visible=False, interactive=True, label="type", scale=1)
 imageboxes.append(img)
 imageboxes.append(typ)
 return imageboxes


'''
define model
'''
config_file = "configs/anycontrol.yaml"
model_file = "ckpts/anycontrol_15.ckpt"
model = create_model(config_file).cpu()
model.load_state_dict(load_state_dict(model_file, location='cuda'))
model = model.cuda()
ddim_sampler = DDIMSampler(model)



block = gr.Blocks(theme='bethecloud/storj_theme').queue()
with block:
 with gr.Row():
 gr.Markdown("## AnyControl Demo")
 gr.Markdown("---")
 with gr.Row():
 with gr.Column(scale=1):
 with gr.Blocks():
 s = gr.Slider(1, max_conditions, value=1, step=1, label="How many conditions to upload:")
 imageboxes = []
 for i in range(max_conditions):
 if i==0:
 with gr.Row():
 img = gr.Image(visible=True, sources=['upload'], type="numpy", label='Condition 1', interactive=True, scale=3, height=200)
 typ = gr.Dropdown(condition_types, visible=True, interactive=True, label="type", scale=1)
 else:
 with gr.Row():
 img = gr.Image(visible=False, sources=['upload'], type="numpy", label=f'Condition {i+1}', scale=3, height=200)
 typ = gr.Dropdown(condition_types, visible=False, interactive=True, label="type", scale=1)
 imageboxes.append(img)
 imageboxes.append(typ)
 s.change(variable_image_outputs, s, imageboxes)
 with gr.Column(scale=2):
 with gr.Row():
 prompt = gr.Textbox(label="Prompt")
 with gr.Row():
 with gr.Column():
 with gr.Accordion("Advanced options", open=False):
 num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=4, step=1)
 image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
 strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1, step=0.01)

 local_strength = gr.Slider(label="Local Strength", minimum=0, maximum=2, value=1, step=0.01)
 global_strength = gr.Slider(label="Global Strength", minimum=0, maximum=2, value=1, step=0.01)
 color_strength = gr.Slider(label="Color Strength", minimum=0, maximum=2, value=1, step=0.01)
 
 ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=50, step=1)
 scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
 seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, value=42, step=1)
 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, extra digit, fewer digits, cropped, worst quality, low quality')


 with gr.Row():
 run_button = gr.Button(value="Run")
 with gr.Row():
 image_gallery = gr.Gallery(label='Generation', show_label=True, elem_id="gallery", columns=[4], rows=[1], height='auto', interactive=False)
 with gr.Row():
 cond_gallery = gr.Gallery(label='Condition', show_label=True, elem_id="gallery", columns=[4], rows=[1], height='auto', interactive=False)

 inputs = [prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, 
 strength, scale, seed, eta, local_strength, global_strength, color_strength] + imageboxes
 run_button.click(fn=process, inputs=inputs, outputs=[image_gallery, cond_gallery])


# uncomment this block in case you need it
# os.environ['http_proxy'] = ''
# os.environ['https_proxy'] = ''
# os.environ['no_proxy'] = 'localhost,127.0.0.0/8,127.0.1.1'
# os.environ['HTTP_PROXY'] = ''
# os.environ['HTTPS_PROXY'] = ''
# os.environ['NO_PROXY'] = 'localhost,127.0.0.0/8,127.0.1.1'
# os.environ['TMPDIR'] = './tmpfiles'


block.launch(server_name='0.0.0.0', allowed_paths=["."], share=False)