Skip to content

stable audio broken with partial loading (indicative of slightly bigger bug) #11855

@doctorpangloss

Description

@doctorpangloss

Custom Node Testing

Expected Behavior

Modules using torch.nn.utils.parametrizations.weight_norm like Conv1d layers in AudioOobleckVAE should be properly loaded to GPU during novram

Actual Behavior

weight-normalized modules are incorrectly skipped in _load_list(), causing:

RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor) should be the same

this affects the Stable Audio VAE and any other model using PyTorch's weight normalization parametrizations.

Steps to Reproduce

  1. Load stable-audio-open-1.0.safetensors
  2. Run audio generation workflow with --novram
  3. VAE decode fails with device mismatch error

Debug Logs

(salient logs shared above)

Other

in stable audio, weight normalization stores weights as parametrizations.weight.original0 and parametrizations.weight.original1, which only appear in recurse=True so:

  for name, param in m.named_parameters(recurse=True):
      if name not in params:
          if not name.startswith("parametrizations."):
              skip = True
              break
parametrizations

is the intended misspelling and is correct

Metadata

Metadata

Assignees

No one assigned

    Labels

    BugSomething is confirmed to not be working properly.

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions