From f8a3fd9d79837bd377d4e15e634271b488d9ee26 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Sat, 25 Jul 2026 06:34:40 +1000 Subject: [PATCH] upscalers: convert latent_upsampler model to DynamicVram (#15063) These were alll non-dynamic (some non-ModelPatcher) code path calling FreeMemory for management requiring up-front memory freeing. Convert it to dynamic to avoid legacy free behaviour mixing into otherwise dynamic workflows. --- comfy/ldm/lightricks/latent_upsampler.py | 35 ++++++++++++++---------- comfy_extras/nodes_hunyuan.py | 9 ++++-- comfy_extras/nodes_lt_upsampler.py | 24 ++++++---------- comfy_extras/nodes_upscale_model.py | 35 +++++++++++------------- 4 files changed, 52 insertions(+), 51 deletions(-) diff --git a/comfy/ldm/lightricks/latent_upsampler.py b/comfy/ldm/lightricks/latent_upsampler.py index 78ed7653f..6a4beb1bf 100644 --- a/comfy/ldm/lightricks/latent_upsampler.py +++ b/comfy/ldm/lightricks/latent_upsampler.py @@ -97,11 +97,11 @@ class SpatialRationalResampler(nn.Module): For dims==3, work per-frame for spatial scaling (temporal axis untouched). """ - def __init__(self, mid_channels: int, scale: float): + def __init__(self, mid_channels: int, scale: float, operations): super().__init__() self.scale = float(scale) self.num, self.den = _rational_for_scale(self.scale) - self.conv = nn.Conv2d( + self.conv = operations.Conv2d( mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1 ) self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) @@ -119,18 +119,18 @@ class SpatialRationalResampler(nn.Module): class ResBlock(nn.Module): def __init__( - self, channels: int, mid_channels: Optional[int] = None, dims: int = 3 + self, channels: int, operations, mid_channels: Optional[int] = None, dims: int = 3 ): super().__init__() if mid_channels is None: mid_channels = channels - Conv = nn.Conv2d if dims == 2 else nn.Conv3d + Conv = operations.Conv2d if dims == 2 else operations.Conv3d self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) - self.norm1 = nn.GroupNorm(32, mid_channels) + self.norm1 = operations.GroupNorm(32, mid_channels) self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) - self.norm2 = nn.GroupNorm(32, channels) + self.norm2 = operations.GroupNorm(32, channels) self.activation = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -159,6 +159,7 @@ class LatentUpsampler(nn.Module): def __init__( self, + operations, in_channels: int = 128, mid_channels: int = 512, num_blocks_per_stage: int = 4, @@ -179,34 +180,34 @@ class LatentUpsampler(nn.Module): self.spatial_scale = float(spatial_scale) self.rational_resampler = rational_resampler - Conv = nn.Conv2d if dims == 2 else nn.Conv3d + Conv = operations.Conv2d if dims == 2 else operations.Conv3d self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1) - self.initial_norm = nn.GroupNorm(32, mid_channels) + self.initial_norm = operations.GroupNorm(32, mid_channels) self.initial_activation = nn.SiLU() self.res_blocks = nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)] ) if spatial_upsample and temporal_upsample: self.upsampler = nn.Sequential( - nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + operations.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(3), ) elif spatial_upsample: if rational_resampler: self.upsampler = SpatialRationalResampler( - mid_channels=mid_channels, scale=self.spatial_scale + mid_channels=mid_channels, scale=self.spatial_scale, operations=operations ) else: self.upsampler = nn.Sequential( - nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + operations.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(2), ) elif temporal_upsample: self.upsampler = nn.Sequential( - nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + operations.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(1), ) else: @@ -215,11 +216,14 @@ class LatentUpsampler(nn.Module): ) self.post_upsample_res_blocks = nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)] ) self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1) + def get_dtype(self): + return getattr(self.initial_conv, "weight_comfy_model_dtype", self.initial_conv.weight.dtype) + def forward(self, latent: torch.Tensor) -> torch.Tensor: b, c, f, h, w = latent.shape @@ -266,7 +270,7 @@ class LatentUpsampler(nn.Module): return x @classmethod - def from_config(cls, config): + def from_config(cls, config, operations): return cls( in_channels=config.get("in_channels", 4), mid_channels=config.get("mid_channels", 128), @@ -276,6 +280,7 @@ class LatentUpsampler(nn.Module): temporal_upsample=config.get("temporal_upsample", False), spatial_scale=config.get("spatial_scale", 2.0), rational_resampler=config.get("rational_resampler", False), + operations=operations, ) def config(self): diff --git a/comfy_extras/nodes_hunyuan.py b/comfy_extras/nodes_hunyuan.py index 8df2c8908..ce2997245 100644 --- a/comfy_extras/nodes_hunyuan.py +++ b/comfy_extras/nodes_hunyuan.py @@ -2,6 +2,8 @@ import nodes import node_helpers import torch import comfy.model_management +import comfy.model_patcher +import comfy.ops from typing_extensions import override from comfy_api.latest import ComfyExtension, io from comfy.ldm.hunyuan_video.upsampler import HunyuanVideo15SRModel @@ -217,8 +219,11 @@ class LatentUpscaleModelLoader(io.ComfyNode): model.load_sd(sd) elif "post_upsample_res_blocks.0.conv2.bias" in sd: config = json.loads(metadata["config"]) - model = LatentUpsampler.from_config(config).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) - model.load_state_dict(sd) + model = LatentUpsampler.from_config(config, operations=comfy.ops.disable_weight_init).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) + comfy.model_management.archive_model_dtypes(model) + model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device()) + model.load_state_dict(sd, assign=model_patcher.is_dynamic()) + model = model_patcher return io.NodeOutput(model) diff --git a/comfy_extras/nodes_lt_upsampler.py b/comfy_extras/nodes_lt_upsampler.py index ef36109d1..7e7975495 100644 --- a/comfy_extras/nodes_lt_upsampler.py +++ b/comfy_extras/nodes_lt_upsampler.py @@ -38,26 +38,20 @@ class LTXVLatentUpsampler(IO.ComfyNode): Returns: tuple: Tuple containing the upsampled latent """ - device = model_management.get_torch_device() - memory_required = model_management.module_size(upscale_model) - - model_dtype = next(upscale_model.parameters()).dtype + device = upscale_model.load_device + model = upscale_model.model + model_dtype = upscale_model.model_dtype() latents = samples["samples"] input_dtype = latents.dtype - memory_required += math.prod(latents.shape) * 3000.0 # TODO: more accurate - model_management.free_memory(memory_required, device) + memory_required = math.prod(latents.shape) * 3000.0 # TODO: more accurate + model_management.load_models_gpu([upscale_model], memory_required=memory_required) - try: - upscale_model.to(device) # TODO: use the comfy model management system. + latents = latents.to(dtype=model_dtype, device=device) - latents = latents.to(dtype=model_dtype, device=device) - - """Upsample latents without tiling.""" - latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents) - upsampled_latents = upscale_model(latents) - finally: - upscale_model.cpu() + """Upsample latents without tiling.""" + latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents) + upsampled_latents = model(latents) upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize( upsampled_latents diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index 1cf5a5d01..a4d692955 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -7,6 +7,7 @@ import folder_paths from typing_extensions import override from comfy_api.latest import ComfyExtension, io import comfy.model_management +import comfy.model_patcher try: from spandrel_extra_arches import EXTRA_REGISTRY @@ -42,6 +43,7 @@ class UpscaleModelLoader(io.ComfyNode): if not isinstance(out, ImageModelDescriptor): raise Exception("Upscale model must be a single-image model.") + out.patcher = comfy.model_patcher.CoreModelPatcher(out.model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device()) return io.NodeOutput(out) load_model = execute # TODO: remove @@ -66,14 +68,12 @@ class ImageUpscaleWithModel(io.ComfyNode): @classmethod def execute(cls, upscale_model, image) -> io.NodeOutput: - device = model_management.get_torch_device() + device = upscale_model.patcher.load_device - memory_required = model_management.module_size(upscale_model.model) - memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate + memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate memory_required += image.nelement() * image.element_size() - model_management.free_memory(memory_required, device) + model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required) - upscale_model.to(device) in_img = image.movedim(-1,-3).to(device) tile = 512 @@ -82,20 +82,17 @@ class ImageUpscaleWithModel(io.ComfyNode): output_device = comfy.model_management.intermediate_device() oom = True - try: - while oom: - try: - steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) - pbar = comfy.utils.ProgressBar(steps) - s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device) - oom = False - except Exception as e: - model_management.raise_non_oom(e) - tile //= 2 - if tile < 128: - raise e - finally: - upscale_model.to("cpu") + while oom: + try: + steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) + pbar = comfy.utils.ProgressBar(steps) + s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device) + oom = False + except Exception as e: + model_management.raise_non_oom(e) + tile //= 2 + if tile < 128: + raise e s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype()) return io.NodeOutput(s)