From 9cd96dc7e7edbf39fc5a357deae22b02cbb383c0 Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Fri, 29 May 2026 11:34:10 +0800
Subject: [PATCH] wan

---
 backend/diffusion_engine/wan.py |  7 +++++++
 modules/sd_models.py            | 10 ++++++++++
 2 files changed, 17 insertions(+)

diff --git a/backend/diffusion_engine/wan.py b/backend/diffusion_engine/wan.py
index 747cab9e..beccc33f 100644
--- a/backend/diffusion_engine/wan.py
+++ b/backend/diffusion_engine/wan.py
@@ -92,6 +92,13 @@ class Wan(ForgeDiffusionEngine):
                 end_image = self.end_image.movedim(1, -1)
                 _, h, w, _ = end_image.shape
 
+        if self.start_image is not None and self.end_image is not None:
+            memory_management.logger.info("[Wan] FirstLastFrameToVideo")
+        elif self.start_image is not None and self.end_image is None:
+            memory_management.logger.info("[Wan] ImageToVideo")
+        elif self.start_image is None and self.end_image is not None:
+            memory_management.logger.info("[Wan] LastFrameToVideo")
+
         image = torch.ones((length, h, w, 3), device="cpu", dtype=torch.float32).mul(0.5)
         mask = torch.ones((1, 1, latent_shape[2] * 4, latent_shape[-2], latent_shape[-1]), device="cpu", dtype=torch.float32)
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
index d5259965..4e0e4943 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -329,7 +329,14 @@ def forge_model_reload():
 
     timer = Timer()
 
+    _last_frame: torch.Tensor = None
+    # maintain the end_image so that FirstLastFrame is not broken
+    # when reloading Wan 2.2 models via Refiner
+
     if model_data.sd_model is not None:
+        if getattr(model_data.sd_model, "end_image", None) is not None:
+            _last_frame = model_data.sd_model.end_image.clone()
+
         model_data.sd_model = None
         model_data.forge_hash = ""
         memory_management.unload_all_models()
@@ -370,6 +377,9 @@ def forge_model_reload():
     sd_model.sd_model_hash = checkpoint_info.calculate_shorthash()
     timer.record("calculate hash")
 
+    if _last_frame is not None:
+        setattr(sd_model, "end_image", _last_frame)
+
     shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
     model_data.set_sd_model(sd_model)