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)