CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

LastFrame for Wan

Browse Source
This commit is contained in:
Haoming 2026-04-30 14:16:19 +07:00 committed by GitHub
parent ed6c6b7ec2
commit 68e8bd5fbd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 109 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

108
backend/diffusion_engine/wan.py
View File

@ -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)

35
extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
View File

@ -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 <b>Flux.1-Kontext</b> / <b>Flux.2-Klein</b> / <b>Qwen-Image-Edit</b> ; Use in <b>txt2img</b> to achieve the effect of empty latent with custom resolution<br>
For <b>Flux.1-Kontext</b> / <b>Flux.2-Klein</b> / <b>Qwen-Image-Edit</b>: Use in <b>txt2img</b> to achieve the effect of empty latent with custom resolution<br>
For <b>Wan 2.2 I2V</b>: Use in <b>txt2img</b> to set as the Last Frame to achieve LastFrameToVideo<br>
<b>Note:</b> This doesn't actually stitch the images ; <b>Tip:</b> Use the "Image to Upload" to paste images
"""
i2i_info = """
For <b>Flux.1-Kontext</b> / <b>Flux.2-Klein</b> / <b>Qwen-Image-Edit</b> ; Use in <b>img2img</b> to achieve the effect of multiple input images<br>
For <b>Flux.1-Kontext</b> / <b>Flux.2-Klein</b> / <b>Qwen-Image-Edit</b>: Use in <b>img2img</b> to achieve the effect of multiple input images<br>
For <b>Wan 2.2 I2V</b>: Use in <b>img2img</b> to set as the Last Frame to achieve FirstLastFrameToVideo<br>
<b>Note:</b> This doesn't actually stitch the images ; <b>Tip:</b> 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