mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-12 23:17:30 +08:00
Compare commits
1 Commits
cloud-open
...
pid_15
| Author | SHA1 | Date | |
|---|---|---|---|
| 18396717b7 |
@ -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:
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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."),
|
||||
|
||||
45
openapi.yaml
45
openapi.yaml
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
Reference in New Issue
Block a user