diff --git a/backend/args.py b/backend/args.py index 917f90e3..5057795e 100644 --- a/backend/args.py +++ b/backend/args.py @@ -162,10 +162,14 @@ class dynamic_args(metaclass=_DynamicArgsMeta): """Flux.2 Klein""" wan: bool = False """Wan 2.2""" + pid: bool = False + """PiD""" ref_latents: list["torch.Tensor"] = [] """Reference Latent(s) for Flux Kontext / Qwen-Image-Edit / Flux.2 Klein""" concat_latent: "torch.Tensor" = None """Input Latent for Wan 2.2 I2V""" + lq_latent: tuple["torch.Tensor", "torch.Tensor"] = None + """lq_latent & degrade_sigma for PiD""" is_referencing: bool = False """Appending Reference Latent(s) (by. ImageStitch)""" ops: str = None diff --git a/backend/loader.py b/backend/loader.py index 9d841f3f..701219d8 100644 --- a/backend/loader.py +++ b/backend/loader.py @@ -867,13 +867,14 @@ def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = No except AttributeError: raise ValueError("Failed to recognize model...") from None - repo_name = estimated_config.huggingface_repo + repo_name: str = estimated_config.huggingface_repo backend.args.dynamic_args.kontext = "kontext" in str(sd).lower() backend.args.dynamic_args.edit = "qwen" in str(sd).lower() and "edit" in str(sd).lower() backend.args.dynamic_args.nunchaku = getattr(estimated_config, "nunchaku", False) backend.args.dynamic_args.klein = "klein" in repo_name backend.args.dynamic_args.wan = "Wan" in repo_name + backend.args.dynamic_args.pid = "PiD" in repo_name if "xl" in repo_name and "rectified" in str(sd).lower(): estimated_config.sampling_settings["RF"] = True diff --git a/backend/nn/pixeldit/pid.py b/backend/nn/pixeldit/pid.py index 6161f7af..c9321b68 100644 --- a/backend/nn/pixeldit/pid.py +++ b/backend/nn/pixeldit/pid.py @@ -4,6 +4,8 @@ import torch import torch.nn as nn import torch.nn.functional as F +from backend.args import dynamic_args + from .model import PixDiT_T2I from .modules import precompute_freqs_cis_2d @@ -178,7 +180,8 @@ class PidNet(PixDiT_T2I): return s return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx) - def forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs): + def forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs): + lq_latent, degrade_sigma = dynamic_args.lq_latent assert lq_latent is not None expected_c = self.lq_proj.latent_channels