Add block replace transformer_options to flux.

comfy/ldm/flux/model.py

@@ -96,7 +96,9 @@ class Flux(nn.Module):
         y: Tensor,
         guidance: Tensor = None,
         control=None,
+        transformer_options={},
     ) -> Tensor:
+        patches_replace = transformer_options.get("patches_replace", {})
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
 
@@ -114,7 +116,18 @@ class Flux(nn.Module):
         ids = torch.cat((txt_ids, img_ids), dim=1)
         pe = self.pe_embedder(ids)
 
+        blocks_replace = patches_replace.get("dit", {})
         for i, block in enumerate(self.double_blocks):
+            if ("double_block", i) in blocks_replace:
+                def block_wrap(args):
+                    out = {}
+                    out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"])
+                    return out
+
+                out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe}, {"original_block": block_wrap})
+                txt = out["txt"]
+                img = out["img"]
+            else:
                 img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
 
             if control is not None: # Controlnet
@@ -127,6 +140,15 @@ class Flux(nn.Module):
         img = torch.cat((txt, img), 1)
 
         for i, block in enumerate(self.single_blocks):
+            if ("single_block", i) in blocks_replace:
+                def block_wrap(args):
+                    out = {}
+                    out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"])
+                    return out
+
+                out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe}, {"original_block": block_wrap})
+                img = out["img"]
+            else:
                 img = block(img, vec=vec, pe=pe)
 
             if control is not None: # Controlnet
@@ -141,7 +163,7 @@ class Flux(nn.Module):
         img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def forward(self, x, timestep, context, y, guidance, control=None, **kwargs):
+    def forward(self, x, timestep, context, y, guidance, control=None, transformer_options={}, **kwargs):
         bs, c, h, w = x.shape
         patch_size = 2
         x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
@@ -156,5 +178,5 @@ class Flux(nn.Module):
         img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
 
         txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
-        out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control)
+        out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options)
         return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]

comfy_extras/nodes_sd3.py

@@ -128,6 +128,9 @@ class SkipLayerGuidanceSD3:
         sigma_start = model_sampling.percent_to_sigma(start_percent)
         sigma_end = model_sampling.percent_to_sigma(end_percent)
 
+        layers = re.findall(r'\d+', layers)
+        layers = [int(i) for i in layers]
+
         def post_cfg_function(args):
             model = args["model"]
             cond_pred = args["cond_denoised"]
@@ -147,8 +150,6 @@ class SkipLayerGuidanceSD3:
                 cfg_result = cfg_result + (cond_pred - slg) * scale
             return cfg_result
 
-        layers = re.findall(r'\d+', layers)
-        layers = [int(i) for i in layers]
         m = model.clone()
         m.set_model_sampler_post_cfg_function(post_cfg_function)