From 68e8bd5fbdbe7efe7568519e1e7452f388fee142 Mon Sep 17 00:00:00 2001
From: Haoming <73768377+Haoming02@users.noreply.github.com>
Date: Thu, 30 Apr 2026 14:16:19 +0700
Subject: [PATCH] LastFrame for Wan
---
backend/diffusion_engine/wan.py | 108 +++++++++++++-----
.../scripts/image_stitch.py | 35 ++++--
2 files changed, 109 insertions(+), 34 deletions(-)
diff --git a/backend/diffusion_engine/wan.py b/backend/diffusion_engine/wan.py
index 5d02b8af..747cab9e 100644
--- a/backend/diffusion_engine/wan.py
+++ b/backend/diffusion_engine/wan.py
@@ -1,8 +1,10 @@
import torch
from huggingface_guess import model_list
-from backend import args, memory_management
+from backend import memory_management
+from backend.args import dynamic_args
from backend.diffusion_engine.base import ForgeDiffusionEngine, ForgeObjects
+from backend.misc.image_resize import adaptive_resize
from backend.modules.k_prediction import PredictionDiscreteFlow
from backend.patcher.clip import CLIP
from backend.patcher.unet import UnetPatcher
@@ -43,15 +45,28 @@ class Wan(ForgeDiffusionEngine):
global refiner_shift
if refiner_shift is not None:
- self.forge_objects.unet.model.predictor.set_parameters(shift=refiner_shift)
- memory_management.logger.debug(f"Shift: {refiner_shift}")
+ super().set_shift(refiner_shift)
refiner_shift = None
+ del self.ini_latent
+ del self.ref_latents
+
+ self.start_image: torch.Tensor = None
+ """first frame; cleared automatically every generation"""
+ self.end_image: torch.Tensor = None
+ """last frame; cleared manually by ImageStitch"""
+
def set_shift(self, shift):
global refiner_shift
super().set_shift(shift)
refiner_shift = shift
+ def clear_references(self):
+ # called by ImageStitch
+ self.start_image = None
+ self.end_image = None
+ memory_management.soft_empty_cache()
+
@torch.inference_mode()
def get_learned_conditioning(self, prompt: list[str]):
memory_management.load_model_gpu(self.forge_objects.clip.patcher)
@@ -63,23 +78,45 @@ class Wan(ForgeDiffusionEngine):
return token_count, max(510, token_count)
@torch.inference_mode()
- def image_to_video(self, length: int, start_image: torch.Tensor, noise: torch.Tensor):
- _, h, w, c = start_image.shape
+ def image_to_video(self, length: int, latent_shape: list[int]):
+ # https://github.com/Comfy-Org/ComfyUI/blob/v0.20.1/comfy_extras/nodes_wan.py#L209
- _image = torch.ones((length, h, w, c), device=start_image.device, dtype=start_image.dtype) * 0.5
- _image[: start_image.shape[0]] = start_image
+ if self.start_image is not None:
+ start_image = self.start_image.movedim(1, -1)
+ _, h, w, _ = start_image.shape
- concat_latent_image = self.forge_objects.vae.encode(_image[:, :, :, :3])
- mask = torch.ones((1, 1, noise.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
- mask[:, :, : ((start_image.shape[0] - 1) // 4) + 1] = 0.0
+ if self.end_image is not None:
+ if self.start_image is not None:
+ end_image = adaptive_resize(self.end_image, w, h, "bilinear", "center").movedim(1, -1)
+ else:
+ end_image = self.end_image.movedim(1, -1)
+ _, h, w, _ = end_image.shape
- image = concat_latent_image
+ 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)
- extra_channels = self.forge_objects.unet.model.diffusion_model.in_dim - 16 # 20
+ if self.start_image is not None:
+ image[: start_image.shape[0]] = start_image
+ mask[:, :, : start_image.shape[0] + 3] = 0.0
- for i in range(0, image.shape[1], 16):
- image[:, i : i + 16] = self.forge_objects.vae.first_stage_model.process_in(image[:, i : i + 16])
- image = resize_to_batch_size(image, noise.shape[0])
+ if self.end_image is not None:
+ image[-end_image.shape[0] :] = end_image
+ mask[:, :, -end_image.shape[0] :] = 0.0
+
+ concat_latent_image = self.forge_objects.vae.encode(image[:, :, :, :3])
+ concat_mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2)
+
+ # https://github.com/Comfy-Org/ComfyUI/blob/v0.20.1/comfy/model_base.py#L1291
+
+ image: torch.Tensor = concat_latent_image
+ mask: torch.Tensor = concat_mask
+
+ extra_channels: int = 20
+ latent_dim: int = 16
+
+ for i in range(0, image.shape[1], latent_dim):
+ image[:, i : i + latent_dim] = self.forge_objects.vae.first_stage_model.process_in(image[:, i : i + latent_dim])
+ image = resize_to_batch_size(image, latent_shape[0])
if image.shape[1] > (extra_channels - 4):
image = image[:, : (extra_channels - 4)]
@@ -87,30 +124,47 @@ class Wan(ForgeDiffusionEngine):
if mask.shape[1] != 4:
mask = torch.mean(mask, dim=1, keepdim=True)
mask = (1.0 - mask).to(image)
- if mask.shape[-3] < noise.shape[-3]:
- mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode="constant", value=0)
+ mask = adaptive_resize(mask, latent_shape[-1], latent_shape[-2], "bilinear", "center")
+ if mask.shape[-3] < latent_shape[-3]:
+ mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, latent_shape[-3] - mask.shape[-3]), mode="constant", value=0)
if mask.shape[1] == 1:
mask = mask.repeat(1, 4, 1, 1, 1)
- mask = resize_to_batch_size(mask, noise.shape[0])
+ mask = resize_to_batch_size(mask, latent_shape[0])
- _concat_mask_index = 0 # TODO
+ z = torch.cat((mask, image), dim=1)
- if _concat_mask_index != 0:
- z = torch.cat((image[:, :_concat_mask_index], mask, image[:, _concat_mask_index:]), dim=1)
- else:
- z = torch.cat((mask, image), dim=1)
+ dynamic_args.concat_latent = z.cpu()
- args.dynamic_args.concat_latent = z
+ self.start_image = None
@torch.inference_mode()
def encode_first_stage(self, x: torch.Tensor):
- b, c, h, w = x.shape
+ b, _, h, w = x.shape
if x.size(0) > 1:
x = x[0].unsqueeze(0) # enforce batch_size of 1
+ x = x.mul(0.5).add(0.5)
+
+ if dynamic_args.is_referencing:
+ if b == 1:
+ # FirstLastFrameToVideo
+ self.end_image = x.cpu()
+ return
+ else:
+ # LastFrameToVideo
+ self.end_image = x.cpu()
+
+ else:
+ if b == 1:
+ # img2img
+ sample = self.forge_objects.vae.encode(x.movedim(1, -1))
+ sample = self.forge_objects.vae.first_stage_model.process_in(sample)
+ return sample.to(x)
+ else:
+ # FirstFrameToVideo
+ self.start_image = x.cpu()
- start_image = x.movedim(1, -1) * 0.5 + 0.5
latent = torch.zeros([1, 16, ((b - 1) // 4) + 1, h // 8, w // 8], device=self.forge_objects.vae.device)
- self.image_to_video(b, start_image, latent)
+ self.image_to_video(b, list(latent.shape))
sample = self.forge_objects.vae.first_stage_model.process_in(latent)
return sample.to(x)
diff --git a/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py b/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
index 9dca32fd..5f50b122 100644
--- a/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
+++ b/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
@@ -6,18 +6,20 @@ from PIL import Image
from backend.args import dynamic_args
from modules import images, scripts, sd_models
from modules.api import api
-from modules.processing import StableDiffusionProcessing
+from modules.processing import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, logger
from modules.sd_samplers_common import images_tensor_to_samples
from modules.shared import device, opts
from modules.ui_components import FormRow, InputAccordion
t2i_info = """
-For Flux.1-Kontext / Flux.2-Klein / Qwen-Image-Edit ; Use in txt2img to achieve the effect of empty latent with custom resolution
+For Flux.1-Kontext / Flux.2-Klein / Qwen-Image-Edit: Use in txt2img to achieve the effect of empty latent with custom resolution
+For Wan 2.2 I2V: Use in txt2img to set as the Last Frame to achieve LastFrameToVideo
Note: This doesn't actually stitch the images ; Tip: Use the "Image to Upload" to paste images
"""
i2i_info = """
-For Flux.1-Kontext / Flux.2-Klein / Qwen-Image-Edit ; Use in img2img to achieve the effect of multiple input images
+For Flux.1-Kontext / Flux.2-Klein / Qwen-Image-Edit: Use in img2img to achieve the effect of multiple input images
+For Wan 2.2 I2V: Use in img2img to set as the Last Frame to achieve FirstLastFrameToVideo
Note: This doesn't actually stitch the images ; Tip: Use the "Image to Upload" to paste images
"""
@@ -151,7 +153,7 @@ class ImageStitch(scripts.Script):
p.sd_model.clear_references()
def process(self, p: StableDiffusionProcessing, enable: bool, references: list[str | tuple[Image.Image, str]], max_dim: int):
- if not (enable and references and any(getattr(dynamic_args, key) for key in ("kontext", "edit", "klein"))):
+ if not (enable and references and any(getattr(dynamic_args, key) for key in ("kontext", "edit", "klein", "wan"))):
if ImageStitch.cached_parameters is None:
return
@@ -163,23 +165,42 @@ class ImageStitch(scripts.Script):
references = self.extract_images(references)
# cache is based on reference inputs & model
- cache: list[str | int] = [str(sd_models.model_data.forge_loading_parameters), *(self.hash_image(ref) for ref in references)]
+ cache: list[str | int | bool] = [str(sd_models.model_data.forge_loading_parameters), *(self.hash_image(ref) for ref in references), (dynamic_args.wan and isinstance(p, StableDiffusionProcessingTxt2Img))]
if ImageStitch.cached_parameters == cache:
return
ImageStitch.cached_parameters = cache
self.reset_references(p)
+ _batch_size: int = None
+
+ if dynamic_args.wan:
+ if isinstance(p, StableDiffusionProcessingTxt2Img):
+ _batch_size = p.batch_size
+ if _batch_size == 1:
+ logger.error("Wan 2.2 requires more than one frame...")
+ return
+ if len(references) > 1:
+ logger.warning("Wan 2.2 only uses the first reference image...")
+ references = [references[0]]
+
dynamic_args.is_referencing = True
for reference in references:
reference = self.preprocess(reference, max_dim)
+ if _batch_size:
+ reference = images.resize_image(1, reference, p.width, p.height)
image = images.flatten(reference, opts.img2img_background_color)
image = np.array(image, dtype=np.float32) / 255.0
image = np.moveaxis(image, 2, 0)
- image = torch.from_numpy(image).to(device=device, dtype=torch.float32)
+ image = torch.from_numpy(image).to(device=device).unsqueeze(0)
- images_tensor_to_samples(image.unsqueeze(0), 0, p.sd_model) # calls encode_first_stage
+ if _batch_size:
+ dim = [_batch_size - 1] + list(image.shape)[1:]
+ empty = torch.empty(dim, dtype=torch.float32, device=device)
+ image = torch.cat([image, empty], dim=0)
+
+ images_tensor_to_samples(image, 0, p.sd_model) # calls encode_first_stage
dynamic_args.is_referencing = False