update app.py

app.py

@@ -1,7 +1,269 @@
+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
 
-def greet(name):
- return "Hello " + name + "!!"
+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
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+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)