Add seg adapter

model.py

@@ -177,7 +177,8 @@ class Model:
         base_model_file = "https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt"
         base_model_file_anything = "https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.0-pruned.ckpt"
         sketch_adapter_file = "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_sketch_sd14v1.pth"
-        pose_adapter_file = "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_keypose_sd14v1.pth" 
+        pose_adapter_file = "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_keypose_sd14v1.pth"
+        seg_adapter_file = "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_seg_sd14v1.pth"
         pidinet_file = model_path+"table5_pidinet.pth"
         clip_file = "https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/*"
         
@@ -185,6 +186,7 @@ class Model:
         subprocess.run(shlex.split(f'wget {base_model_file_anything} -O models/anything-v4.0-pruned.ckpt'))
         subprocess.run(shlex.split(f'wget {sketch_adapter_file} -O models/t2iadapter_sketch_sd14v1.pth'))
         subprocess.run(shlex.split(f'wget {pose_adapter_file} -O models/t2iadapter_keypose_sd14v1.pth'))
+        subprocess.run(shlex.split(f'wget {seg_adapter_file} -O models/t2iadapter_seg_sd14v1.pth'))
         subprocess.run(shlex.split(f'wget {pidinet_file} -O models/table5_pidinet.pth'))
 
         
@@ -204,6 +206,9 @@ class Model:
         self.model_ad_pose = Adapter(cin=int(3*64),channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device)
         self.model_ad_pose.load_state_dict(torch.load("models/t2iadapter_keypose_sd14v1.pth"))
 
+        self.model_ad_seg = Adapter(cin=int(3*64),channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device)
+        self.model_ad_seg.load_state_dict(torch.load("models/t2iadapter_seg_sd14v1.pth""))
+
 
     @torch.inference_mode()
     def process_sketch(self, input_img, type_in, color_back, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):    
@@ -370,4 +375,66 @@ class Model:
         x_samples_ddim = 255.*x_samples_ddim
         x_samples_ddim = x_samples_ddim.astype(np.uint8)
     
-        return [im_pose[:,:,::-1].astype(np.uint8), x_samples_ddim]
+        return [im_pose[:,:,::-1].astype(np.uint8), x_samples_ddim]
+
+
+    @torch.inference_mode()
+    def process_seg(self, input_img, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):    
+        global current_base
+        device = 'cuda' 
+        if base_model == 'sd-v1-4.ckpt':
+            model = self.model
+            sampler = self.sampler
+        else:
+            model = self.model_anything
+            sampler = self.sampler_anything
+        # if current_base != base_model:
+        #     ckpt = os.path.join("models", base_model)
+        #     pl_sd = torch.load(ckpt, map_location="cpu")
+        #     if "state_dict" in pl_sd:
+        #         sd = pl_sd["state_dict"]
+        #     else:
+        #         sd = pl_sd
+        #     model.load_state_dict(sd, strict=False) #load_model_from_config(config, os.path.join("models", base_model)).to(device)
+        #     current_base = base_model
+        con_strength = int((1-con_strength)*50)
+        if fix_sample == 'True':
+            seed_everything(42)
+        
+        im = cv2.resize(input_img,(512,512))
+        mask = im.copy()
+        mask = img2tensor(mask, bgr2rgb=True, float32=True)/255.
+        mask = mask.unsqueeze(0)
+
+        im_mask = tensor2img(mask)
+
+        c = model.get_learned_conditioning([prompt])
+        nc = model.get_learned_conditioning([neg_prompt])
+        
+        with torch.no_grad():
+            # extract condition features
+            features_adapter = self.model_ad_seg(mask.to(device))
+    
+        shape = [4, 64, 64]
+    
+        # sampling
+        samples_ddim, _ = sampler.sample(S=50,
+                                        conditioning=c,
+                                        batch_size=1,
+                                        shape=shape,
+                                        verbose=False,
+                                        unconditional_guidance_scale=scale,
+                                        unconditional_conditioning=nc,
+                                        eta=0.0,
+                                        x_T=None,
+                                        features_adapter1=features_adapter,
+                                        mode = 'mask',
+                                        con_strength = con_strength)
+    
+        x_samples_ddim = model.decode_first_stage(samples_ddim)
+        x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
+        x_samples_ddim = x_samples_ddim.permute(0, 2, 3, 1).cpu().numpy()[0]
+        x_samples_ddim = 255.*x_samples_ddim
+        x_samples_ddim = x_samples_ddim.astype(np.uint8)
+    
+        return [im_edge, x_samples_ddim]