From 7f6dfdbfe3337f3ec93600c5f5606c13047af92a Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Thu, 12 Feb 2026 17:06:49 +0800
Subject: [PATCH] batch

---
 backend/args.py       | 2 ++
 backend/loader.py     | 1 +
 modules/processing.py | 9 +++++----
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/backend/args.py b/backend/args.py
index da896d29..c8920bec 100644
--- a/backend/args.py
+++ b/backend/args.py
@@ -145,6 +145,7 @@ dynamic_args = dict(
     edit=False,
     nunchaku=False,
     klein=False,
+    wan=False,
     ref_latents=[],
     concat_latent=None,
     is_referencing=False,
@@ -159,6 +160,7 @@ Some parameters that are used throughout the Webui
 - edit: `bool` - Qwen-Image-Edit
 - nunchaku: `bool` - Nunchaku (SVDQ) Models
 - klein: `bool` - Flux.2 Klein
+- wan: `bool` - Wan 2.2
 - ref_latents: `list[torch.Tensor]` - Reference Latent(s) for Flux Kontext & Qwen-Image-Edit
 - concat_latent: `torch.Tensor` - Input Latent for Wan 2.2 I2V
 - is_referencing: `bool` - Appending Reference Latent(s) (by. ImageStitch)
diff --git a/backend/loader.py b/backend/loader.py
index 5307e232..ead99660 100644
--- a/backend/loader.py
+++ b/backend/loader.py
@@ -676,6 +676,7 @@ def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = No
     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
 
     if getattr(estimated_config, "nunchaku", False):
         estimated_config.unet_config["filename"] = str(sd)
diff --git a/modules/processing.py b/modules/processing.py
index 50b500f4..7024fb2e 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -22,7 +22,7 @@ import modules.paths as paths
 import modules.sd_models as sd_models
 import modules.sd_vae as sd_vae
 import modules.shared as shared
-from backend import memory_management
+from backend import args, memory_management
 from backend.modules.k_prediction import rescale_zero_terminal_snr_sigmas
 from backend.utils import hash_tensor
 from modules import devices, errors, extra_networks, images, infotext_utils, masking, profiling, prompt_parser, rng, scripts, sd_samplers, sd_samplers_common, sd_unet, sd_vae_approx
@@ -849,9 +849,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
 
+    _times = 1
     _is_video = False
     video_path = None
-    if shared.sd_model.is_wan:
+    if shared.sd_model.is_wan and args.dynamic_args["wan"]:
         _times = ((getattr(p, "batch_size", 1) - 1) // 4) + 1  # https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy_extras/nodes_wan.py#L41
         p.batch_size = (_times - 1) * 4 + 1
         _is_video: bool = _times > 1
@@ -1386,7 +1387,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             # here we generate an image normally
 
             x = self.rng.next()
-            if shared.sd_model.is_wan:  # enforce batch_size of 1
+            if shared.sd_model.is_wan and args.dynamic_args["wan"]:  # enforce batch_size of 1
                 x = x[0].unsqueeze(0)
 
             self.sd_model.forge_objects = self.sd_model.forge_objects_after_applying_lora.shallow_copy()
@@ -1858,7 +1859,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 
     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
         x = self.rng.next()
-        if shared.sd_model.is_wan:  # enforce batch_size of 1
+        if shared.sd_model.is_wan and args.dynamic_args["wan"]:  # enforce batch_size of 1
             x = x[0].unsqueeze(0)
 
         if self.initial_noise_multiplier != 1.0: