Refactor cond_concat into model object.

comfy/samplers.py

@@ -358,15 +358,6 @@ def sgm_scheduler(model, steps):
     sigs += [0.0]
     return torch.FloatTensor(sigs)
 
-def blank_inpaint_image_like(latent_image):
-    blank_image = torch.ones_like(latent_image)
-    # these are the values for "zero" in pixel space translated to latent space
-    blank_image[:,0] *= 0.8223
-    blank_image[:,1] *= -0.6876
-    blank_image[:,2] *= 0.6364
-    blank_image[:,3] *= 0.1380
-    return blank_image
-
 def get_mask_aabb(masks):
     if masks.numel() == 0:
         return torch.zeros((0, 4), device=masks.device, dtype=torch.int)
@@ -671,21 +662,10 @@ def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model
 
     extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
 
-    cond_concat = None
-    if hasattr(model, 'concat_keys'): #inpaint
-        cond_concat = []
-        for ck in model.concat_keys:
-            if denoise_mask is not None:
-                if ck == "mask":
-                    cond_concat.append(denoise_mask[:,:1])
-                elif ck == "masked_image":
-                    cond_concat.append(latent_image) #NOTE: the latent_image should be masked by the mask in pixel space
-            else:
-                if ck == "mask":
-                    cond_concat.append(torch.ones_like(noise)[:,:1])
-                elif ck == "masked_image":
-                    cond_concat.append(blank_inpaint_image_like(noise))
-        extra_args["cond_concat"] = cond_concat
+    if hasattr(model, 'cond_concat'):
+        cond_concat = model.cond_concat(noise=noise, latent_image=latent_image, denoise_mask=denoise_mask)
+        if cond_concat is not None:
+            extra_args["cond_concat"] = cond_concat
 
     samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
     return model.process_latent_out(samples.to(torch.float32))