add free-hand and anything

app.py

@@ -18,10 +18,14 @@ from torch import autocast
 from ldm.inference_base import (DEFAULT_NEGATIVE_PROMPT, diffusion_inference, get_adapters, get_sd_models)
 from ldm.modules.extra_condition import api
 from ldm.modules.extra_condition.api import (ExtraCondition, get_adapter_feature, get_cond_model)
+import numpy as np
+from ldm.util import read_state_dict
 
 torch.set_grad_enabled(False)
 
-supported_cond = ['style', 'color', 'sketch', 'openpose', 'depth', 'canny']
+supported_cond_map = ['style', 'color', 'sketch', 'openpose', 'depth', 'canny']
+supported_cond = ['style', 'color', 'sketch', 'sketch', 'openpose', 'depth', 'canny']
+draw_map = gr.Interface(lambda x: x, gr.Image(source="canvas"), gr.Image())
 
 # download the checkpoints
 urls = {
@@ -35,6 +39,7 @@ urls = {
         "models/t2iadapter_sketch_sd15v2.pth"
     ],
     'runwayml/stable-diffusion-v1-5': ['v1-5-pruned-emaonly.ckpt'],
+    'CompVis/stable-diffusion-v-1-4-original':['sd-v1-4.ckpt'],
     'andite/anything-v4.0': ['anything-v4.0-pruned.ckpt', 'anything-v4.0.vae.pt'],
 }
 
@@ -93,93 +98,130 @@ global_opt.sampler = 'ddim'
 global_opt.cond_weight = 1.0
 global_opt.C = 4
 global_opt.f = 8
-
-# stable-diffusion model
-sd_model, sampler = get_sd_models(global_opt)
 # adapters and models to processing condition inputs
 adapters = {}
 cond_models = {}
 torch.cuda.empty_cache()
 
 
-def run(*args):
-    with torch.inference_mode(), \
-            sd_model.ema_scope(), \
-            autocast('cuda'):
-
-        inps = []
-        for i in range(0, len(args) - 8, len(supported_cond)):
-            inps.append(args[i:i + len(supported_cond)])
+def draw_transfer(im1):
+    c = im1[:, :, 0:3].astype(np.float32)
+    a = im1[:, :, 3:4].astype(np.float32) / 255.0
+    im1 = c * a + 255.0 * (1.0 - a)
+    im1 = (im1.clip(0, 255)).astype(np.uint8)
 
-        opt = copy.deepcopy(global_opt)
-        opt.prompt, opt.neg_prompt, opt.scale, opt.n_samples, opt.seed, opt.steps, opt.resize_short_edge, opt.cond_tau \
-            = args[-8:]
+    return im1
 
-        conds = []
-        activated_conds = []
+class process:
+    def __init__(self):
+        self.base_model = 'v1-5-pruned-emaonly.ckpt'
+        # stable-diffusion model
+        self.sd_model, self.sampler = get_sd_models(global_opt)
 
-        ims1 = []
-        ims2 = []
-        for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):
-            if idx > 0:
+    def run(self, *args):
+        opt = copy.deepcopy(global_opt)
+        opt.prompt, opt.neg_prompt, opt.scale, opt.n_samples, opt.seed, opt.steps, opt.resize_short_edge, opt.cond_tau, opt.base_model \
+                = args[-9:]
+        # check base model
+        if opt.base_model!=self.base_model:
+            ckpt = os.path.join("models", opt.base_model)
+            pl_sd = read_state_dict(ckpt)
+            if "state_dict" in pl_sd:
+                st = pl_sd["state_dict"]
+            else:
+                st = pl_sd
+            self.sd_model.load_state_dict(st, strict=False)
+            self.base_model = opt.base_model
+            if self.base_model!='v1-5-pruned-emaonly.ckpt' and self.base_model!='sd-v1-4.ckpt':
+                vae_sd = torch.load(os.path.join('models', 'anything-v4.0.vae.pt'), map_location="cuda")
+                st = vae_sd["state_dict"]
+                self.sd_model.first_stage_model.load_state_dict(st, strict=False)
+
+        with torch.inference_mode(), \
+                self.sd_model.ema_scope(), \
+                autocast('cuda'):
+
+            inps = []
+            for i in range(0, len(args) - 9, len(supported_cond)):
+                inps.append(args[i:i + len(supported_cond)])
+
+            conds = []
+            activated_conds = []
+
+            ims1 = []
+            ims2 = []
+            for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):
                 if b != 'Nothing' and (im1 is not None or im2 is not None):
+                    if im1 is not None and isinstance(im1,dict):
+                        im1 = im1['mask']
+                        im1 = draw_transfer(im1)
+
                     if im1 is not None:
                         h, w, _ = im1.shape
                     else:
                         h, w, _ = im2.shape
-                    # break
-        # resize all the images to the same size
-        for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):
-            if idx == 0:
+
+            # resize all the images to the same size
+            for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):
+                if idx == 0:
+                    ims1.append(im1)
+                    ims2.append(im2)
+                    continue
+                if b != 'Nothing':
+                    if im1 is not None and isinstance(im1,dict):
+                            im1 = im1['mask']
+                            im1 = draw_transfer(im1)
+                            im2 = im1
+                            cv2.imwrite('sketch.png', im1)
+                    if im1 is not None:
+                        im1 = cv2.resize(im1, (w, h), interpolation=cv2.INTER_CUBIC)
+                    if im2 is not None:
+                        im2 = cv2.resize(im2, (w, h), interpolation=cv2.INTER_CUBIC)
                 ims1.append(im1)
                 ims2.append(im2)
-                continue
-            if b != 'Nothing':
-                if im1 is not None:
-                    im1 = cv2.resize(im1, (w, h), interpolation=cv2.INTER_CUBIC)
-                if im2 is not None:
-                    im2 = cv2.resize(im2, (w, h), interpolation=cv2.INTER_CUBIC)
-            ims1.append(im1)
-            ims2.append(im2)
-
-        for idx, (b, _, _, cond_weight) in enumerate(zip(*inps)):
-            cond_name = supported_cond[idx]
-            if b == 'Nothing':
-                if cond_name in adapters:
-                    adapters[cond_name]['model'] = adapters[cond_name]['model'].cpu()
-            else:
-                activated_conds.append(cond_name)
-                if cond_name in adapters:
-                    adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)
+
+            for idx, (b, _, _, cond_weight) in enumerate(zip(*inps)):
+                cond_name = supported_cond[idx]
+                if b == 'Nothing':
+                    if cond_name in adapters:
+                        adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)#.cpu()
                 else:
-                    adapters[cond_name] = get_adapters(opt, getattr(ExtraCondition, cond_name))
-                adapters[cond_name]['cond_weight'] = cond_weight
+                    # print(idx,b)
+                    activated_conds.append(cond_name)
+                    if cond_name in adapters:
+                        adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)
+                    else:
+                        adapters[cond_name] = get_adapters(opt, getattr(ExtraCondition, cond_name))
+                    adapters[cond_name]['cond_weight'] = cond_weight
 
-                process_cond_module = getattr(api, f'get_cond_{cond_name}')
+                    process_cond_module = getattr(api, f'get_cond_{cond_name}')
 
-                if b == 'Image':
-                    if cond_name not in cond_models:
-                        cond_models[cond_name] = get_cond_model(opt, getattr(ExtraCondition, cond_name))
-                    conds.append(process_cond_module(opt, ims1[idx], 'image', cond_models[cond_name]))
-                else:
-                    conds.append(process_cond_module(opt, ims2[idx], cond_name, None))
+                    if b == 'Image':
+                        if cond_name not in cond_models:
+                            cond_models[cond_name] = get_cond_model(opt, getattr(ExtraCondition, cond_name))
+                        conds.append(process_cond_module(opt, ims1[idx], 'image', cond_models[cond_name]))
+                    else:
+                        if idx == 2: # draw
+                            conds.append(process_cond_module(opt, (255.-ims2[idx]).astype(np.uint8), cond_name, None))
+                        else:
+                            conds.append(process_cond_module(opt, ims2[idx], cond_name, None))
 
-        adapter_features, append_to_context = get_adapter_feature(
-            conds, [adapters[cond_name] for cond_name in activated_conds])
+            adapter_features, append_to_context = get_adapter_feature(
+                conds, [adapters[cond_name] for cond_name in activated_conds])
 
-        output_conds = []
-        for cond in conds:
-            output_conds.append(tensor2img(cond, rgb2bgr=False))
+            output_conds = []
+            for cond in conds:
+                output_conds.append(tensor2img(cond, rgb2bgr=False))
 
-        ims = []
-        seed_everything(opt.seed)
-        for _ in range(opt.n_samples):
-            result = diffusion_inference(opt, sd_model, sampler, adapter_features, append_to_context)
-            ims.append(tensor2img(result, rgb2bgr=False))
+            ims = []
+            seed_everything(opt.seed)
+            for _ in range(opt.n_samples):
+                result = diffusion_inference(opt, self.sd_model, self.sampler, adapter_features, append_to_context)
+                ims.append(tensor2img(result, rgb2bgr=False))
 
-        # Clear GPU memory cache so less likely to OOM
-        torch.cuda.empty_cache()
-        return ims, output_conds
+            # Clear GPU memory cache so less likely to OOM
+            torch.cuda.empty_cache()
+            return ims, output_conds
 
 
 def change_visible(im1, im2, val):
@@ -195,13 +237,14 @@ def change_visible(im1, im2, val):
         outputs[im2] = gr.update(visible=True)
     return outputs
 
-
-DESCRIPTION = '# [Composable T2I-Adapter](https://github.com/TencentARC/T2I-Adapter)'
+DESCRIPTION = '# [T2I-Adapter](https://github.com/TencentARC/T2I-Adapter)'
 
 DESCRIPTION += f'<p>Gradio demo for **T2I-Adapter**: [[GitHub]](https://github.com/TencentARC/T2I-Adapter), [[Paper]](https://arxiv.org/abs/2302.08453). If T2I-Adapter is helpful, please help to ⭐ the [Github Repo](https://github.com/TencentARC/T2I-Adapter) and recommend it to your friends 😊 </p>'
 
 DESCRIPTION += f'<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/Adapter/T2I-Adapter?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'
 
+processer = process()
+
 with gr.Blocks(css='style.css') as demo:
     gr.Markdown(DESCRIPTION)
 
@@ -215,7 +258,7 @@ with gr.Blocks(css='style.css') as demo:
             with gr.Box():
                 gr.Markdown("<h5><center>Style & Color</center></h5>")
                 with gr.Row():
-                    for cond_name in supported_cond[:2]:
+                    for cond_name in supported_cond_map[:2]:
                         with gr.Box():
                             with gr.Column():
                                 if cond_name == 'style':
@@ -232,6 +275,7 @@ with gr.Blocks(css='style.css') as demo:
                                         interactive=True,
                                         value="Nothing",
                                     )
+
                                 im1 = gr.Image(
                                     source='upload', label="Image", interactive=True, visible=False, type="numpy")
                                 im2 = gr.Image(
@@ -251,11 +295,38 @@ with gr.Blocks(css='style.css') as demo:
                                 ims1.append(im1)
                                 ims2.append(im2)
                                 cond_weights.append(cond_weight)
+
+            with gr.Box():
+                gr.Markdown("<h5><center>Drawing</center></h5>")
+                with gr.Column():
+                    btn1 = gr.Radio(
+                        choices=["Sketch", "Nothing"],
+                        label=f"Input type for drawing",
+                        interactive=True,
+                        value="Nothing")
+                    im1 = gr.Image(source='canvas', tool='color-sketch', label='Pay attention to adjusting stylus thickness!', visible=False) 
+                    im2 = im1
+                    cond_weight = gr.Slider(
+                        label="Condition weight",
+                        minimum=0,
+                        maximum=5,
+                        step=0.05,
+                        value=1,
+                        interactive=True)
+
+                    fn = partial(change_visible, im1, im2)
+                    btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)
+
+                    btns.append(btn1)
+                    ims1.append(im1)
+                    ims2.append(im2)
+                    cond_weights.append(cond_weight)
+
         with gr.Column(scale=4):
             with gr.Box():
                 gr.Markdown("<h5><center>Structure</center></h5>")
                 with gr.Row():
-                    for cond_name in supported_cond[2:6]:
+                    for cond_name in supported_cond_map[2:6]:
                         with gr.Box():
                             with gr.Column():
                                 if cond_name == 'openpose':
@@ -272,6 +343,7 @@ with gr.Blocks(css='style.css') as demo:
                                         interactive=True,
                                         value="Nothing",
                                     )
+
                                 im1 = gr.Image(
                                     source='upload', label="Image", interactive=True, visible=False, type="numpy")
                                 im2 = gr.Image(
@@ -286,44 +358,46 @@ with gr.Blocks(css='style.css') as demo:
 
                                 fn = partial(change_visible, im1, im2)
                                 btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)
-
                                 btns.append(btn1)
                                 ims1.append(im1)
                                 ims2.append(im2)
                                 cond_weights.append(cond_weight)
 
-    with gr.Column():
-        prompt = gr.Textbox(label="Prompt")
-
-        with gr.Accordion('Advanced options', open=False):
-            neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
-            scale = gr.Slider(
-                label="Guidance Scale (Classifier free guidance)", value=7.5, minimum=1, maximum=20, step=0.1)
-            n_samples = gr.Slider(label="Num samples", value=1, minimum=1, maximum=1, step=1)
-            seed = gr.Slider(label="Seed", value=42, minimum=0, maximum=10000, step=1, randomize=True)
-            steps = gr.Slider(label="Steps", value=50, minimum=10, maximum=100, step=1)
-            resize_short_edge = gr.Slider(label="Image resolution", value=512, minimum=320, maximum=1024, step=1)
-            cond_tau = gr.Slider(
-                label="timestamp parameter that determines until which step the adapter is applied",
-                value=1.0,
-                minimum=0.1,
-                maximum=1.0,
-                step=0.05)
-
-    with gr.Row():
-        submit = gr.Button("Generate")
-    output = gr.Gallery().style(grid=2, height='auto')
-    cond = gr.Gallery().style(grid=2, height='auto')
+            with gr.Column():
+                base_model = gr.inputs.Radio(['v1-5-pruned-emaonly.ckpt', 'sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='v1-5-pruned-emaonly.ckpt', label='The base model you want to use. You can try more base models on https://civitai.com/.')
+                prompt = gr.Textbox(label="Prompt")
+                with gr.Accordion('Advanced options', open=False):
+                    neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
+                    scale = gr.Slider(
+                        label="Guidance Scale (Classifier free guidance)", value=7.5, minimum=1, maximum=20, step=0.1)
+                    n_samples = gr.Slider(label="Num samples", value=1, minimum=1, maximum=1, step=1)
+                    seed = gr.Slider(label="Seed", value=42, minimum=0, maximum=10000, step=1, randomize=True)
+                    steps = gr.Slider(label="Steps", value=50, minimum=10, maximum=100, step=1)
+                    resize_short_edge = gr.Slider(label="Image resolution", value=512, minimum=320, maximum=1024, step=1)
+                    cond_tau = gr.Slider(
+                        label="timestamp parameter that determines until which step the adapter is applied",
+                        value=1.0,
+                        minimum=0.1,
+                        maximum=1.0,
+                        step=0.05)
+                submit = gr.Button("Generate")
+
+    with gr.Box():
+        gr.Markdown("<h5><center>Results</center></h5>")
+        with gr.Column():
+            output = gr.Gallery().style(grid=2, height='auto')
+            cond = gr.Gallery().style(grid=2, height='auto')
 
     inps = list(chain(btns, ims1, ims2, cond_weights))
 
-    inps.extend([prompt, neg_prompt, scale, n_samples, seed, steps, resize_short_edge, cond_tau])
-    submit.click(fn=run, inputs=inps, outputs=[output, cond])
+    inps.extend([prompt, neg_prompt, scale, n_samples, seed, steps, resize_short_edge, cond_tau, base_model])
+    submit.click(fn=processer.run, inputs=inps, outputs=[output, cond])
 
     ex = gr.Examples([
         [
             "Image",
             "Nothing",
+            "Nothing",
             "Image",
             "Nothing",
             "Nothing",
@@ -346,6 +420,7 @@ with gr.Blocks(css='style.css') as demo:
             1,
             1,
             1,
+            1,
             "master sword",
             "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
             7.5,
@@ -358,6 +433,7 @@ with gr.Blocks(css='style.css') as demo:
         [
             "Image",
             "Nothing",
+            "Nothing",
             "Image",
             "Nothing",
             "Nothing",
@@ -380,6 +456,7 @@ with gr.Blocks(css='style.css') as demo:
             1,
             1,
             1,
+            1,
             "motorcycle",
             "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
             7.5,