RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
rattus 2026-07-25 06:34:40 +10:00 committed by GitHub
parent 7c59a078d6
commit f8a3fd9d79
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 52 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
comfy/ldm/lightricks/latent_upsampler.py
View File

@ -97,11 +97,11 @@ class SpatialRationalResampler(nn.Module):
For dims==3, work per-frame for spatial scaling (temporal axis untouched). 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__() super().__init__()
self.scale = float(scale) self.scale = float(scale)
self.num, self.den = _rational_for_scale(self.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 mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1
) )
self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
@ -119,18 +119,18 @@ class SpatialRationalResampler(nn.Module):
class ResBlock(nn.Module): class ResBlock(nn.Module):
def __init__( 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__() super().__init__()
if mid_channels is None: if mid_channels is None:
mid_channels = channels 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.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.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() self.activation = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
@ -159,6 +159,7 @@ class LatentUpsampler(nn.Module):
def __init__( def __init__(
self, self,
operations,
in_channels: int = 128, in_channels: int = 128,
mid_channels: int = 512, mid_channels: int = 512,
num_blocks_per_stage: int = 4, num_blocks_per_stage: int = 4,
@ -179,34 +180,34 @@ class LatentUpsampler(nn.Module):
self.spatial_scale = float(spatial_scale) self.spatial_scale = float(spatial_scale)
self.rational_resampler = rational_resampler 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_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.initial_activation = nn.SiLU()
self.res_blocks = nn.ModuleList( 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: if spatial_upsample and temporal_upsample:
self.upsampler = nn.Sequential( 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), PixelShuffleND(3),
) )
elif spatial_upsample: elif spatial_upsample:
if rational_resampler: if rational_resampler:
self.upsampler = SpatialRationalResampler( self.upsampler = SpatialRationalResampler(
mid_channels=mid_channels, scale=self.spatial_scale mid_channels=mid_channels, scale=self.spatial_scale, operations=operations
) )
else: else:
self.upsampler = nn.Sequential( 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), PixelShuffleND(2),
) )
elif temporal_upsample: elif temporal_upsample:
self.upsampler = nn.Sequential( 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), PixelShuffleND(1),
) )
else: else:
@ -215,11 +216,14 @@ class LatentUpsampler(nn.Module):
) )
self.post_upsample_res_blocks = nn.ModuleList( 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) 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: def forward(self, latent: torch.Tensor) -> torch.Tensor:
b, c, f, h, w = latent.shape b, c, f, h, w = latent.shape
@ -266,7 +270,7 @@ class LatentUpsampler(nn.Module):
return x return x
@classmethod @classmethod
def from_config(cls, config): def from_config(cls, config, operations):
return cls( return cls(
in_channels=config.get("in_channels", 4), in_channels=config.get("in_channels", 4),
mid_channels=config.get("mid_channels", 128), mid_channels=config.get("mid_channels", 128),
@ -276,6 +280,7 @@ class LatentUpsampler(nn.Module):
temporal_upsample=config.get("temporal_upsample", False), temporal_upsample=config.get("temporal_upsample", False),
spatial_scale=config.get("spatial_scale", 2.0), spatial_scale=config.get("spatial_scale", 2.0),
rational_resampler=config.get("rational_resampler", False), rational_resampler=config.get("rational_resampler", False),
operations=operations,
) )
def config(self): def config(self):

9
comfy_extras/nodes_hunyuan.py
View File

@ -2,6 +2,8 @@ import nodes
import node_helpers import node_helpers
import torch import torch
import comfy.model_management import comfy.model_management
import comfy.model_patcher
import comfy.ops
from typing_extensions import override from typing_extensions import override
from comfy_api.latest import ComfyExtension, io from comfy_api.latest import ComfyExtension, io
from comfy.ldm.hunyuan_video.upsampler import HunyuanVideo15SRModel from comfy.ldm.hunyuan_video.upsampler import HunyuanVideo15SRModel
@ -217,8 +219,11 @@ class LatentUpscaleModelLoader(io.ComfyNode):
model.load_sd(sd) model.load_sd(sd)
elif "post_upsample_res_blocks.0.conv2.bias" in sd: elif "post_upsample_res_blocks.0.conv2.bias" in sd:
config = json.loads(metadata["config"]) config = json.loads(metadata["config"])
model = LatentUpsampler.from_config(config).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) model = LatentUpsampler.from_config(config, operations=comfy.ops.disable_weight_init).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32]))
model.load_state_dict(sd) 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) return io.NodeOutput(model)

24
comfy_extras/nodes_lt_upsampler.py
View File

@ -38,26 +38,20 @@ class LTXVLatentUpsampler(IO.ComfyNode):
Returns: Returns:
tuple: Tuple containing the upsampled latent tuple: Tuple containing the upsampled latent
""" """
device = model_management.get_torch_device() device = upscale_model.load_device
memory_required = model_management.module_size(upscale_model) model = upscale_model.model
model_dtype = upscale_model.model_dtype()
model_dtype = next(upscale_model.parameters()).dtype
latents = samples["samples"] latents = samples["samples"]
input_dtype = latents.dtype input_dtype = latents.dtype
memory_required += math.prod(latents.shape) * 3000.0 # TODO: more accurate memory_required = math.prod(latents.shape) * 3000.0 # TODO: more accurate
model_management.free_memory(memory_required, device) model_management.load_models_gpu([upscale_model], memory_required=memory_required)
try: latents = latents.to(dtype=model_dtype, device=device)
upscale_model.to(device) # TODO: use the comfy model management system.
latents = latents.to(dtype=model_dtype, device=device) """Upsample latents without tiling."""
latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents)
"""Upsample latents without tiling.""" upsampled_latents = model(latents)
latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents)
upsampled_latents = upscale_model(latents)
finally:
upscale_model.cpu()
upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize( upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize(
upsampled_latents upsampled_latents

35
comfy_extras/nodes_upscale_model.py
View File

@ -7,6 +7,7 @@ import folder_paths
from typing_extensions import override from typing_extensions import override
from comfy_api.latest import ComfyExtension, io from comfy_api.latest import ComfyExtension, io
import comfy.model_management import comfy.model_management
import comfy.model_patcher
try: try:
from spandrel_extra_arches import EXTRA_REGISTRY from spandrel_extra_arches import EXTRA_REGISTRY
@ -42,6 +43,7 @@ class UpscaleModelLoader(io.ComfyNode):
if not isinstance(out, ImageModelDescriptor): if not isinstance(out, ImageModelDescriptor):
raise Exception("Upscale model must be a single-image model.") 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) return io.NodeOutput(out)
load_model = execute # TODO: remove load_model = execute # TODO: remove
@ -66,14 +68,12 @@ class ImageUpscaleWithModel(io.ComfyNode):
@classmethod @classmethod
def execute(cls, upscale_model, image) -> io.NodeOutput: 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() 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) in_img = image.movedim(-1,-3).to(device)
tile = 512 tile = 512
@ -82,20 +82,17 @@ class ImageUpscaleWithModel(io.ComfyNode):
output_device = comfy.model_management.intermediate_device() output_device = comfy.model_management.intermediate_device()
oom = True oom = True
try: while oom:
while oom: try:
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)
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)
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)
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
oom = False except Exception as e:
except Exception as e: model_management.raise_non_oom(e)
model_management.raise_non_oom(e) tile //= 2
tile //= 2 if tile < 128:
if tile < 128: raise e
raise e
finally:
upscale_model.to("cpu")
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype()) s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype())
return io.NodeOutput(s) return io.NodeOutput(s)