From a72d152b0c8788b954004ee8c01e957701514e32 Mon Sep 17 00:00:00 2001 From: Bratzmeister Date: Tue, 12 Nov 2024 11:53:36 +0000 Subject: [PATCH 1/4] fix --cuda-device arg for AMD/HIP devices (#5586) * fix --cuda-device arg for AMD/HIP devices CUDA_VISIBLE_DEVICES is ignored for HIP devices/backend. Instead it uses HIP_VISIBLE_DEVICES. Setting this environment variable has no side effect for CUDA/NVIDIA so it can safely be set in any case and vice versa. * deleted accidental if --- main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/main.py b/main.py index c23210861..05eb31c7a 100644 --- a/main.py +++ b/main.py @@ -71,6 +71,7 @@ if os.name == "nt": if __name__ == "__main__": if args.cuda_device is not None: os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) + os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) logging.info("Set cuda device to: {}".format(args.cuda_device)) if args.deterministic: From 8ebf2d8831f3c1ec68f203c12ba6af456f4d99e4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 12 Nov 2024 08:00:00 -0500 Subject: [PATCH 2/4] Add block replace transformer_options to flux. --- comfy/ldm/flux/model.py | 30 ++++++++++++++++++++++++++---- comfy_extras/nodes_sd3.py | 5 +++-- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index 233d78394..ae1ed109d 100644 --- a/comfy/ldm/flux/model.py +++ b/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,8 +116,19 @@ 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): - img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + 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 control_i = control.get("input") @@ -127,7 +140,16 @@ class Flux(nn.Module): img = torch.cat((txt, img), 1) for i, block in enumerate(self.single_blocks): - img = block(img, vec=vec, pe=pe) + 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 control_o = control.get("output") @@ -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] diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py index bbdedef79..e95f20b98 100644 --- a/comfy_extras/nodes_sd3.py +++ b/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) From 3748e7ef7a42197755f30a9b68f3e9b86d59e354 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 13 Nov 2024 04:24:48 -0500 Subject: [PATCH 3/4] Fix regression. --- comfy_extras/nodes_model_advanced.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_model_advanced.py b/comfy_extras/nodes_model_advanced.py index ed14b61ac..f085bf12f 100644 --- a/comfy_extras/nodes_model_advanced.py +++ b/comfy_extras/nodes_model_advanced.py @@ -26,8 +26,8 @@ class X0(comfy.model_sampling.EPS): class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete): original_timesteps = 50 - def __init__(self, model_config=None): - super().__init__(model_config) + def __init__(self, model_config=None, zsnr=None): + super().__init__(model_config, zsnr=zsnr) self.skip_steps = self.num_timesteps // self.original_timesteps From 3b9a6cf2b11094f92228b121c6a0d466ba5d5246 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 13 Nov 2024 07:18:30 -0500 Subject: [PATCH 4/4] Fix issue with 3d masks. --- comfy/utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 04926c1e4..985cd9a1b 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -853,19 +853,18 @@ def reshape_mask(input_mask, output_shape): dims = len(output_shape) - 2 if dims == 1: - mask = input_mask scale_mode = "linear" if dims == 2: - mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1])) + input_mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1])) scale_mode = "bilinear" if dims == 3: if len(input_mask.shape) < 5: - mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1])) + input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1])) scale_mode = "trilinear" - mask = torch.nn.functional.interpolate(mask, size=output_shape[2:], mode=scale_mode) + mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode) if mask.shape[1] < output_shape[1]: mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]] mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0])