Compare commits

..

1 Commits

Author SHA1 Message Date
18396717b7 Support PID 1.5 models. 2026-07-12 00:01:53 -04:00
6 changed files with 162 additions and 42 deletions

View File

@ -197,6 +197,9 @@ class PixDiT_T2I(nn.Module):
"""Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
return s
def _pre_pixel_blocks(self, s, **kwargs):
return s
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
H_orig, W_orig = x.shape[2], x.shape[3]
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
@ -226,6 +229,7 @@ class PixDiT_T2I(nn.Module):
s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
s = F.silu(t_emb + s)
s = self._pre_pixel_blocks(s, **kwargs)
s_cond = s.view(B * L, self.hidden_size)
x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
for blk in self.pixel_blocks:

View File

@ -13,15 +13,15 @@ from .model import PixDiT_T2I
from .modules import precompute_freqs_cis_2d
class SigmaAwareGatePerTokenPerDim(nn.Module):
class SigmaAwareGate(nn.Module):
"""gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
"""
def __init__(self, dim: int, dtype=None, device=None, operations=None):
def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
super().__init__()
self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device)
self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
@ -36,15 +36,15 @@ class SigmaAwareGatePerTokenPerDim(nn.Module):
class ResBlock(nn.Module):
"""Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None):
def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None):
super().__init__()
self.block = nn.Sequential(
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
nn.SiLU(),
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
nn.SiLU(),
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
@ -62,9 +62,13 @@ class LQProjection2D(nn.Module):
patch_size: int = 16,
sr_scale: int = 4,
latent_spatial_down_factor: int = 8,
latent_unpatchify_factor: int = 1,
num_res_blocks: int = 4,
num_outputs: int = 7,
interval: int = 2,
conv_padding_mode: str = "zeros",
gate_per_token: bool = False,
pit_output: bool = False,
dtype=None, device=None, operations=None,
):
super().__init__()
@ -74,34 +78,38 @@ class LQProjection2D(nn.Module):
self.patch_size = patch_size
self.sr_scale = sr_scale
self.latent_spatial_down_factor = latent_spatial_down_factor
self.latent_unpatchify_factor = latent_unpatchify_factor
self.num_outputs = num_outputs
self.interval = interval
z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor)
effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor
z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size
self.z_to_patch_ratio = z_to_patch_ratio
if z_to_patch_ratio >= 1:
self.latent_fold_factor = 0
latent_proj_in_ch = latent_channels
latent_proj_in_ch = effective_latent_channels
else:
fold_factor = int(1 / z_to_patch_ratio)
assert fold_factor * z_to_patch_ratio == 1.0
self.latent_fold_factor = fold_factor
latent_proj_in_ch = latent_channels * fold_factor * fold_factor
latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
layers = [
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
nn.SiLU(),
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
]
for _ in range(num_res_blocks):
layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations))
self.latent_proj = nn.Sequential(*layers)
self.output_heads = nn.ModuleList(
[operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
)
self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None
self.gate_modules = nn.ModuleList(
[SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
[SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
for _ in range(num_outputs)]
)
@ -115,6 +123,11 @@ class LQProjection2D(nn.Module):
return self.gate_modules[out_idx](x, lq_feature, sigma)
def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
f = self.latent_unpatchify_factor
if f > 1:
B, C, H, W = lq_latent.shape
lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W)
lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f)
B, z_dim = lq_latent.shape[:2]
if self.z_to_patch_ratio >= 1:
if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
@ -134,7 +147,10 @@ class LQProjection2D(nn.Module):
feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
B, C, H, W = feat.shape
tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
return [head(tokens) for head in self.output_heads]
outputs = [head(tokens) for head in self.output_heads]
if self.pit_head is not None:
outputs.append(self.pit_head(tokens))
return outputs
class PidNet(PixDiT_T2I):
@ -148,6 +164,10 @@ class PidNet(PixDiT_T2I):
lq_interval: int = 2,
sr_scale: int = 4,
latent_spatial_down_factor: int = 8,
lq_latent_unpatchify_factor: int = 1,
lq_conv_padding_mode: str = "zeros",
lq_gate_per_token: bool = False,
pit_lq_inject: bool = False,
rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
rope_ref_w: int = 1024,
image_model=None,
@ -165,6 +185,8 @@ class PidNet(PixDiT_T2I):
for blk in self.pixel_blocks:
blk._rope_fn = _pit_rope_fn
self.pit_lq_inject = pit_lq_inject
num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
self.lq_proj = LQProjection2D(
latent_channels=lq_latent_channels,
@ -173,13 +195,20 @@ class PidNet(PixDiT_T2I):
patch_size=self.patch_size,
sr_scale=sr_scale,
latent_spatial_down_factor=latent_spatial_down_factor,
latent_unpatchify_factor=lq_latent_unpatchify_factor,
num_res_blocks=lq_num_res_blocks,
num_outputs=num_lq_outputs,
interval=lq_interval,
conv_padding_mode=lq_conv_padding_mode,
gate_per_token=lq_gate_per_token,
pit_output=pit_lq_inject,
dtype=dtype,
device=device,
operations=operations,
)
self.pit_lq_gate = SigmaAwareGate(
self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations
) if pit_lq_inject else None
def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
return precompute_freqs_cis_2d(
@ -197,6 +226,11 @@ class PidNet(PixDiT_T2I):
return s
return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs):
if pid_pit_lq_feature is None:
return s
return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma)
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
if lq_latent is None:
raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
@ -216,12 +250,14 @@ class PidNet(PixDiT_T2I):
degrade_sigma = degrade_sigma.expand(B).contiguous()
lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None
return super()._forward(
x, timesteps,
context=context, attention_mask=attention_mask,
transformer_options=transformer_options,
pid_lq_features=lq_features,
pid_pit_lq_feature=pit_lq_feature,
pid_degrade_sigma=degrade_sigma,
**kwargs,
)

View File

@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
if _lq_w_key in state_dict_keys:
in_ch = int(state_dict[_lq_w_key].shape[1])
latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
hidden_dim = int(state_dict[_lq_w_key].shape[0])
_gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
num_gates = len({k[len(_gate_prefix):].split('.')[0]
for k in state_dict_keys if k.startswith(_gate_prefix)})
pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys
dit_config = {"image_model": "pid",
"lq_latent_channels": in_ch,
"latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
"lq_hidden_dim": hidden_dim}
if num_gates > 0:
dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
if pid_v1_5:
pid_v1_5_variants = {
16: { # Flux and QwenImage
"lq_latent_channels": 16,
"latent_spatial_down_factor": 8,
"lq_latent_unpatchify_factor": 1,
},
32: { # Flux2 after 2x latent unpatchify
"lq_latent_channels": 128,
"latent_spatial_down_factor": 16,
"lq_latent_unpatchify_factor": 2,
},
}
variant = pid_v1_5_variants.get(latent_proj_in_channels)
if variant is None:
raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")
gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]
dit_config.update(variant)
dit_config.update({
"lq_conv_padding_mode": "replicate",
"lq_gate_per_token": gate_weight.shape[0] == 1,
"pit_lq_inject": True,
"rope_ref_h": 2048,
"rope_ref_w": 2048,
})
else:
dit_config.update({
"lq_latent_channels": latent_proj_in_channels,
"latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,
})
return dit_config
if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I

View File

@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
display_name="Save Video",
category="video",
essentials_category="Basics",
description="Saves the input videos to your ComfyUI output directory.",
description="Saves the input images to your ComfyUI output directory.",
inputs=[
io.Video.Input("video", tooltip="The video to save."),
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),

View File

@ -7,18 +7,18 @@ components:
description: Timestamp when the asset was created
format: date-time
type: string
display_name:
description: Display name of the asset. Mirrors name for backwards compatibility.
nullable: true
type: string
file_path:
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
nullable: true
type: string
hash:
description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$
type: string
loader_path:
description: The value a loader consumes to load this asset. Null when no loader can resolve the file.
nullable: true
type: string
display_name:
description: Human-facing label for the asset. Not unique.
nullable: true
type: string
id:
description: Unique identifier for the asset
format: uuid
@ -144,14 +144,6 @@ components:
AssetUpdated:
description: Response returned when an existing asset is successfully updated.
properties:
display_name:
description: Display name of the asset. Mirrors name for backwards compatibility.
nullable: true
type: string
file_path:
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
nullable: true
type: string
hash:
description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$
@ -783,6 +775,14 @@ components:
ModelFolder:
description: Represents a folder containing models
properties:
extensions:
description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all).
example:
- .ckpt
- .safetensors
items:
type: string
type: array
folders:
description: List of paths where models of this type are stored
example:
@ -1644,7 +1644,7 @@ paths:
format: uuid
type: string
tags:
description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order.
description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:<folder_name>` tag. Extra tags are stored as labels and do not create path components.
type: string
user_metadata:
description: Custom JSON metadata as a string
@ -1829,7 +1829,7 @@ paths:
content:
application/json:
schema:
$ref: '#/components/schemas/AssetUpdated'
$ref: '#/components/schemas/Asset'
description: Asset updated successfully
"400":
content:
@ -2470,6 +2470,9 @@ paths:
supports_preview_metadata:
description: Whether the server supports preview metadata
type: boolean
supports_model_type_tags:
description: Whether the server supports namespaced model type asset tags
type: boolean
type: object
description: Success
headers:
@ -3297,12 +3300,6 @@ paths:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Invalid request parameters
"401":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Unauthorized - Authentication required
"500":
content:
application/json:

View File

@ -97,6 +97,21 @@ def _make_seedvr2_3b_shared_mm_sd():
}
def _make_pid_v1_5_sd(latent_proj_channels=16):
sd = {
"pixel_embedder.proj.weight": torch.empty(16, 3, device="meta"),
"lq_proj.latent_proj.0.weight": torch.empty(1024, latent_proj_channels, 3, 3, device="meta"),
"lq_proj.pit_head.weight": torch.empty(1536, 1024, device="meta"),
"lq_proj.gate_modules.0.content_proj.weight": torch.empty(1, 3072, device="meta"),
"pixel_blocks.0.attn.q_norm.weight": torch.empty(72, device="meta"),
"pixel_blocks.0.adaLN_modulation.0.weight": torch.empty(24576, 1536, device="meta"),
"pixel_blocks.0.adaLN_modulation.0.bias": torch.empty(24576, device="meta"),
}
for i in range(7):
sd[f"lq_proj.gate_modules.{i}.log_alpha"] = torch.empty((), device="meta")
return sd
def _add_model_diffusion_prefix(sd):
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
@ -206,6 +221,43 @@ class TestModelDetection:
assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2"
def test_pid_v1_5_detection(self):
sd = _make_pid_v1_5_sd()
unet_config = detect_unet_config(sd, "")
assert unet_config == {
"image_model": "pid",
"lq_latent_channels": 16,
"lq_hidden_dim": 1024,
"latent_spatial_down_factor": 8,
"lq_interval": 2,
"lq_latent_unpatchify_factor": 1,
"lq_conv_padding_mode": "replicate",
"lq_gate_per_token": True,
"pit_lq_inject": True,
"rope_ref_h": 2048,
"rope_ref_w": 2048,
}
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "PiD"
def test_pid_v1_5_flux2_detection(self):
unet_config = detect_unet_config(_make_pid_v1_5_sd(latent_proj_channels=32), "")
assert unet_config["lq_latent_channels"] == 128
assert unet_config["latent_spatial_down_factor"] == 16
assert unet_config["lq_latent_unpatchify_factor"] == 2
def test_pid_v1_5_pixel_adaln_conversion(self):
sd = _make_pid_v1_5_sd()
model_config = model_config_from_unet_config(detect_unet_config(sd, ""), sd)
processed = model_config.process_unet_state_dict(sd)
assert processed["pixel_blocks.0.attn.q_norm.weight"].shape == (72,)
assert processed["pixel_blocks.0.adaLN_modulation_msa.weight"].shape == (12288, 1536)
assert processed["pixel_blocks.0.adaLN_modulation_mlp.weight"].shape == (12288, 1536)
assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
def test_unet_config_and_required_keys_combination_is_unique(self):
"""Each model in the registry must have a unique combination of
``unet_config`` and ``required_keys``. If two models share the same