Spaces:
Running
Running
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. | |
import gc | |
import logging | |
import math | |
import os | |
import random | |
import sys | |
import types | |
from contextlib import contextmanager | |
from functools import partial | |
from mmgp import offload | |
import torch | |
import torch.nn as nn | |
import torch.cuda.amp as amp | |
import torch.distributed as dist | |
from tqdm import tqdm | |
from PIL import Image | |
import torchvision.transforms.functional as TF | |
import torch.nn.functional as F | |
from .distributed.fsdp import shard_model | |
from .modules.model import WanModel | |
from .modules.t5 import T5EncoderModel | |
from .modules.vae import WanVAE | |
from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, | |
get_sampling_sigmas, retrieve_timesteps) | |
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler | |
from wan.modules.posemb_layers import get_rotary_pos_embed | |
from .utils.vace_preprocessor import VaceVideoProcessor | |
from wan.utils.basic_flowmatch import FlowMatchScheduler | |
from wan.utils.utils import get_outpainting_frame_location | |
def optimized_scale(positive_flat, negative_flat): | |
# Calculate dot production | |
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) | |
# Squared norm of uncondition | |
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 | |
# st_star = v_cond^T * v_uncond / ||v_uncond||^2 | |
st_star = dot_product / squared_norm | |
return st_star | |
class WanT2V: | |
def __init__( | |
self, | |
config, | |
checkpoint_dir, | |
rank=0, | |
model_filename = None, | |
model_type = None, | |
base_model_type = None, | |
text_encoder_filename = None, | |
quantizeTransformer = False, | |
save_quantized = False, | |
dtype = torch.bfloat16, | |
VAE_dtype = torch.float32, | |
mixed_precision_transformer = False | |
): | |
self.device = torch.device(f"cuda") | |
self.config = config | |
self.rank = rank | |
self.dtype = dtype | |
self.num_train_timesteps = config.num_train_timesteps | |
self.param_dtype = config.param_dtype | |
self.text_encoder = T5EncoderModel( | |
text_len=config.text_len, | |
dtype=config.t5_dtype, | |
device=torch.device('cpu'), | |
checkpoint_path=text_encoder_filename, | |
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), | |
shard_fn= None) | |
self.vae_stride = config.vae_stride | |
self.patch_size = config.patch_size | |
self.vae = WanVAE( | |
vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), dtype= VAE_dtype, | |
device=self.device) | |
logging.info(f"Creating WanModel from {model_filename[-1]}") | |
from mmgp import offload | |
# model_filename = "c:/temp/vace1.3/diffusion_pytorch_model.safetensors" | |
# model_filename = "Vacefusionix_quanto_fp16_int8.safetensors" | |
# model_filename = "c:/temp/t2v/diffusion_pytorch_model-00001-of-00006.safetensors" | |
# config_filename= "c:/temp/t2v/t2v.json" | |
base_config_file = f"configs/{base_model_type}.json" | |
forcedConfigPath = base_config_file if len(model_filename) > 1 else None | |
self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer and not save_quantized, writable_tensors= False, defaultConfigPath=base_config_file , forcedConfigPath= forcedConfigPath) | |
# offload.load_model_data(self.model, "c:/temp/Phantom-Wan-1.3B.pth") | |
# self.model.to(torch.bfloat16) | |
# self.model.cpu() | |
self.model.lock_layers_dtypes(torch.float32 if mixed_precision_transformer else dtype) | |
# dtype = torch.bfloat16 | |
# offload.load_model_data(self.model, "ckpts/Wan14BT2VFusioniX_fp16.safetensors") | |
offload.change_dtype(self.model, dtype, True) | |
# offload.save_model(self.model, "wan2.1_selforcing_fp16.safetensors", config_file_path=base_config_file) | |
# offload.save_model(self.model, "wan2.1_text2video_14B_mbf16.safetensors", config_file_path=base_config_file) | |
# offload.save_model(self.model, "wan2.1_text2video_14B_quanto_mfp16_int8.safetensors", do_quantize=True, config_file_path=base_config_file) | |
self.model.eval().requires_grad_(False) | |
if save_quantized: | |
from wgp import save_quantized_model | |
save_quantized_model(self.model, model_type, model_filename[1 if base_model_type=="fantasy" else 0], dtype, base_config_file) | |
self.sample_neg_prompt = config.sample_neg_prompt | |
if base_model_type in ["vace_14B", "vace_1.3B"]: | |
self.vid_proc = VaceVideoProcessor(downsample=tuple([x * y for x, y in zip(config.vae_stride, self.patch_size)]), | |
min_area=480*832, | |
max_area=480*832, | |
min_fps=config.sample_fps, | |
max_fps=config.sample_fps, | |
zero_start=True, | |
seq_len=32760, | |
keep_last=True) | |
self.adapt_vace_model() | |
def vace_encode_frames(self, frames, ref_images, masks=None, tile_size = 0, overlapped_latents = None): | |
if ref_images is None: | |
ref_images = [None] * len(frames) | |
else: | |
assert len(frames) == len(ref_images) | |
if masks is None: | |
latents = self.vae.encode(frames, tile_size = tile_size) | |
else: | |
inactive = [i * (1 - m) + 0 * m for i, m in zip(frames, masks)] | |
reactive = [i * m + 0 * (1 - m) for i, m in zip(frames, masks)] | |
inactive = self.vae.encode(inactive, tile_size = tile_size) | |
if overlapped_latents != None and False : | |
# inactive[0][:, 0:1] = self.vae.encode([frames[0][:, 0:1]], tile_size = tile_size)[0] # redundant | |
for t in inactive: | |
t[:, 1:overlapped_latents.shape[1] + 1] = overlapped_latents | |
reactive = self.vae.encode(reactive, tile_size = tile_size) | |
latents = [torch.cat((u, c), dim=0) for u, c in zip(inactive, reactive)] | |
cat_latents = [] | |
for latent, refs in zip(latents, ref_images): | |
if refs is not None: | |
if masks is None: | |
ref_latent = self.vae.encode(refs, tile_size = tile_size) | |
else: | |
ref_latent = self.vae.encode(refs, tile_size = tile_size) | |
ref_latent = [torch.cat((u, torch.zeros_like(u)), dim=0) for u in ref_latent] | |
assert all([x.shape[1] == 1 for x in ref_latent]) | |
latent = torch.cat([*ref_latent, latent], dim=1) | |
cat_latents.append(latent) | |
return cat_latents | |
def vace_encode_masks(self, masks, ref_images=None): | |
if ref_images is None: | |
ref_images = [None] * len(masks) | |
else: | |
assert len(masks) == len(ref_images) | |
result_masks = [] | |
for mask, refs in zip(masks, ref_images): | |
c, depth, height, width = mask.shape | |
new_depth = int((depth + 3) // self.vae_stride[0]) # nb latents token without (ref tokens not included) | |
height = 2 * (int(height) // (self.vae_stride[1] * 2)) | |
width = 2 * (int(width) // (self.vae_stride[2] * 2)) | |
# reshape | |
mask = mask[0, :, :, :] | |
mask = mask.view( | |
depth, height, self.vae_stride[1], width, self.vae_stride[1] | |
) # depth, height, 8, width, 8 | |
mask = mask.permute(2, 4, 0, 1, 3) # 8, 8, depth, height, width | |
mask = mask.reshape( | |
self.vae_stride[1] * self.vae_stride[2], depth, height, width | |
) # 8*8, depth, height, width | |
# interpolation | |
mask = F.interpolate(mask.unsqueeze(0), size=(new_depth, height, width), mode='nearest-exact').squeeze(0) | |
if refs is not None: | |
length = len(refs) | |
mask_pad = torch.zeros_like(mask[:, :length, :, :]) | |
mask = torch.cat((mask_pad, mask), dim=1) | |
result_masks.append(mask) | |
return result_masks | |
def vace_latent(self, z, m): | |
return [torch.cat([zz, mm], dim=0) for zz, mm in zip(z, m)] | |
def fit_image_into_canvas(self, ref_img, image_size, canvas_tf_bg, device, fill_max = False, outpainting_dims = None, return_mask = False): | |
from wan.utils.utils import save_image | |
ref_width, ref_height = ref_img.size | |
if (ref_height, ref_width) == image_size and outpainting_dims == None: | |
ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(1) | |
canvas = torch.zeros_like(ref_img) if return_mask else None | |
else: | |
if outpainting_dims != None: | |
final_height, final_width = image_size | |
canvas_height, canvas_width, margin_top, margin_left = get_outpainting_frame_location(final_height, final_width, outpainting_dims, 8) | |
else: | |
canvas_height, canvas_width = image_size | |
scale = min(canvas_height / ref_height, canvas_width / ref_width) | |
new_height = int(ref_height * scale) | |
new_width = int(ref_width * scale) | |
if fill_max and (canvas_height - new_height) < 16: | |
new_height = canvas_height | |
if fill_max and (canvas_width - new_width) < 16: | |
new_width = canvas_width | |
top = (canvas_height - new_height) // 2 | |
left = (canvas_width - new_width) // 2 | |
ref_img = ref_img.resize((new_width, new_height), resample=Image.Resampling.LANCZOS) | |
ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(1) | |
if outpainting_dims != None: | |
canvas = torch.full((3, 1, final_height, final_width), canvas_tf_bg, dtype= torch.float, device=device) # [-1, 1] | |
canvas[:, :, margin_top + top:margin_top + top + new_height, margin_left + left:margin_left + left + new_width] = ref_img | |
else: | |
canvas = torch.full((3, 1, canvas_height, canvas_width), canvas_tf_bg, dtype= torch.float, device=device) # [-1, 1] | |
canvas[:, :, top:top + new_height, left:left + new_width] = ref_img | |
ref_img = canvas | |
canvas = None | |
if return_mask: | |
if outpainting_dims != None: | |
canvas = torch.ones((3, 1, final_height, final_width), dtype= torch.float, device=device) # [-1, 1] | |
canvas[:, :, margin_top + top:margin_top + top + new_height, margin_left + left:margin_left + left + new_width] = 0 | |
else: | |
canvas = torch.ones((3, 1, canvas_height, canvas_width), dtype= torch.float, device=device) # [-1, 1] | |
canvas[:, :, top:top + new_height, left:left + new_width] = 0 | |
canvas = canvas.to(device) | |
return ref_img.to(device), canvas | |
def prepare_source(self, src_video, src_mask, src_ref_images, total_frames, image_size, device, keep_frames= [], start_frame = 0, fit_into_canvas = None, pre_src_video = None, inject_frames = [], outpainting_dims = None, any_background_ref = False): | |
image_sizes = [] | |
trim_video = len(keep_frames) | |
def conv_tensor(t, device): | |
return t.float().div_(127.5).add_(-1).permute(3, 0, 1, 2).to(device) | |
for i, (sub_src_video, sub_src_mask, sub_pre_src_video) in enumerate(zip(src_video, src_mask,pre_src_video)): | |
prepend_count = 0 if sub_pre_src_video == None else sub_pre_src_video.shape[1] | |
num_frames = total_frames - prepend_count | |
num_frames = min(num_frames, trim_video) if trim_video > 0 else num_frames | |
if sub_src_mask is not None and sub_src_video is not None: | |
src_video[i] = conv_tensor(sub_src_video[:num_frames], device) | |
src_mask[i] = conv_tensor(sub_src_mask[:num_frames], device) | |
# src_video is [-1, 1] (at this function output), 0 = inpainting area (in fact 127 in [0, 255]) | |
# src_mask is [-1, 1] (at this function output), 0 = preserve original video (in fact 127 in [0, 255]) and 1 = Inpainting (in fact 255 in [0, 255]) | |
if prepend_count > 0: | |
src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) | |
src_mask[i] = torch.cat( [torch.full_like(sub_pre_src_video, -1.0), src_mask[i]] ,1) | |
src_video_shape = src_video[i].shape | |
if src_video_shape[1] != total_frames: | |
src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) | |
src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) | |
src_mask[i] = torch.clamp((src_mask[i][:1, :, :, :] + 1) / 2, min=0, max=1) | |
image_sizes.append(src_video[i].shape[2:]) | |
elif sub_src_video is None: | |
if prepend_count > 0: | |
src_video[i] = torch.cat( [sub_pre_src_video, torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)], dim=1) | |
src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), torch.ones((3, num_frames, image_size[0], image_size[1]), device=device)] ,1) | |
else: | |
src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device) | |
src_mask[i] = torch.ones_like(src_video[i], device=device) | |
image_sizes.append(image_size) | |
else: | |
src_video[i] = conv_tensor(sub_src_video[:num_frames], device) | |
src_mask[i] = torch.ones_like(src_video[i], device=device) | |
if prepend_count > 0: | |
src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) | |
src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1) | |
src_video_shape = src_video[i].shape | |
if src_video_shape[1] != total_frames: | |
src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) | |
src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) | |
image_sizes.append(src_video[i].shape[2:]) | |
for k, keep in enumerate(keep_frames): | |
if not keep: | |
src_video[i][:, k:k+1] = 0 | |
src_mask[i][:, k:k+1] = 1 | |
for k, frame in enumerate(inject_frames): | |
if frame != None: | |
src_video[i][:, k:k+1], src_mask[i][:, k:k+1] = self.fit_image_into_canvas(frame, image_size, 0, device, True, outpainting_dims, return_mask= True) | |
self.background_mask = None | |
for i, ref_images in enumerate(src_ref_images): | |
if ref_images is not None: | |
image_size = image_sizes[i] | |
for j, ref_img in enumerate(ref_images): | |
if ref_img is not None and not torch.is_tensor(ref_img): | |
if j==0 and any_background_ref: | |
if self.background_mask == None: self.background_mask = [None] * len(src_ref_images) | |
src_ref_images[i][j], self.background_mask[i] = self.fit_image_into_canvas(ref_img, image_size, 0, device, True, outpainting_dims, return_mask= True) | |
else: | |
src_ref_images[i][j], _ = self.fit_image_into_canvas(ref_img, image_size, 1, device) | |
if self.background_mask != None: | |
self.background_mask = [ item if item != None else self.background_mask[0] for item in self.background_mask ] # deplicate background mask with double control net since first controlnet image ref modifed by ref | |
return src_video, src_mask, src_ref_images | |
def decode_latent(self, zs, ref_images=None, tile_size= 0 ): | |
if ref_images is None: | |
ref_images = [None] * len(zs) | |
# else: | |
# assert len(zs) == len(ref_images) | |
trimed_zs = [] | |
for z, refs in zip(zs, ref_images): | |
if refs is not None: | |
z = z[:, len(refs):, :, :] | |
trimed_zs.append(z) | |
return self.vae.decode(trimed_zs, tile_size= tile_size) | |
def get_vae_latents(self, ref_images, device, tile_size= 0): | |
ref_vae_latents = [] | |
for ref_image in ref_images: | |
ref_image = TF.to_tensor(ref_image).sub_(0.5).div_(0.5).to(self.device) | |
img_vae_latent = self.vae.encode([ref_image.unsqueeze(1)], tile_size= tile_size) | |
ref_vae_latents.append(img_vae_latent[0]) | |
return torch.cat(ref_vae_latents, dim=1) | |
def generate(self, | |
input_prompt, | |
input_frames= None, | |
input_masks = None, | |
input_ref_images = None, | |
input_video=None, | |
target_camera=None, | |
context_scale=None, | |
width = 1280, | |
height = 720, | |
fit_into_canvas = True, | |
frame_num=81, | |
shift=5.0, | |
sample_solver='unipc', | |
sampling_steps=50, | |
guide_scale=5.0, | |
n_prompt="", | |
seed=-1, | |
offload_model=True, | |
callback = None, | |
enable_RIFLEx = None, | |
VAE_tile_size = 0, | |
joint_pass = False, | |
slg_layers = None, | |
slg_start = 0.0, | |
slg_end = 1.0, | |
cfg_star_switch = True, | |
cfg_zero_step = 5, | |
overlapped_latents = None, | |
return_latent_slice = None, | |
overlap_noise = 0, | |
conditioning_latents_size = 0, | |
model_filename = None, | |
**bbargs | |
): | |
r""" | |
Generates video frames from text prompt using diffusion process. | |
Args: | |
input_prompt (`str`): | |
Text prompt for content generation | |
size (tupele[`int`], *optional*, defaults to (1280,720)): | |
Controls video resolution, (width,height). | |
frame_num (`int`, *optional*, defaults to 81): | |
How many frames to sample from a video. The number should be 4n+1 | |
shift (`float`, *optional*, defaults to 5.0): | |
Noise schedule shift parameter. Affects temporal dynamics | |
sample_solver (`str`, *optional*, defaults to 'unipc'): | |
Solver used to sample the video. | |
sampling_steps (`int`, *optional*, defaults to 40): | |
Number of diffusion sampling steps. Higher values improve quality but slow generation | |
guide_scale (`float`, *optional*, defaults 5.0): | |
Classifier-free guidance scale. Controls prompt adherence vs. creativity | |
n_prompt (`str`, *optional*, defaults to ""): | |
Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` | |
seed (`int`, *optional*, defaults to -1): | |
Random seed for noise generation. If -1, use random seed. | |
offload_model (`bool`, *optional*, defaults to True): | |
If True, offloads models to CPU during generation to save VRAM | |
Returns: | |
torch.Tensor: | |
Generated video frames tensor. Dimensions: (C, N H, W) where: | |
- C: Color channels (3 for RGB) | |
- N: Number of frames (81) | |
- H: Frame height (from size) | |
- W: Frame width from size) | |
""" | |
# preprocess | |
vace = "Vace" in model_filename | |
if n_prompt == "": | |
n_prompt = self.sample_neg_prompt | |
seed = seed if seed >= 0 else random.randint(0, sys.maxsize) | |
seed_g = torch.Generator(device=self.device) | |
seed_g.manual_seed(seed) | |
if self._interrupt: | |
return None | |
context = self.text_encoder([input_prompt], self.device)[0] | |
context_null = self.text_encoder([n_prompt], self.device)[0] | |
context = context.to(self.dtype) | |
context_null = context_null.to(self.dtype) | |
input_ref_images_neg = None | |
phantom = False | |
if target_camera != None: | |
width = input_video.shape[2] | |
height = input_video.shape[1] | |
input_video = input_video.to(dtype=self.dtype , device=self.device) | |
input_video = input_video.permute(3, 0, 1, 2).div_(127.5).sub_(1.) | |
source_latents = self.vae.encode([input_video])[0] #.to(dtype=self.dtype, device=self.device) | |
del input_video | |
# Process target camera (recammaster) | |
from wan.utils.cammmaster_tools import get_camera_embedding | |
cam_emb = get_camera_embedding(target_camera) | |
cam_emb = cam_emb.to(dtype=self.dtype, device=self.device) | |
if vace : | |
# vace context encode | |
input_frames = [u.to(self.device) for u in input_frames] | |
input_ref_images = [ None if u == None else [v.to(self.device) for v in u] for u in input_ref_images] | |
input_masks = [u.to(self.device) for u in input_masks] | |
if self.background_mask != None: self.background_mask = [m.to(self.device) for m in self.background_mask] | |
previous_latents = None | |
# if overlapped_latents != None: | |
# input_ref_images = [u[-1:] for u in input_ref_images] | |
z0 = self.vace_encode_frames(input_frames, input_ref_images, masks=input_masks, tile_size = VAE_tile_size, overlapped_latents = overlapped_latents ) | |
m0 = self.vace_encode_masks(input_masks, input_ref_images) | |
if self.background_mask != None: | |
zbg = self.vace_encode_frames([ref_img[0] for ref_img in input_ref_images], None, masks=self.background_mask, tile_size = VAE_tile_size ) | |
mbg = self.vace_encode_masks(self.background_mask, None) | |
for zz0, mm0, zzbg, mmbg in zip(z0, m0, zbg, mbg): | |
zz0[:, 0:1] = zzbg | |
mm0[:, 0:1] = mmbg | |
self.background_mask = zz0 = mm0 = zzbg = mmbg = None | |
z = self.vace_latent(z0, m0) | |
target_shape = list(z0[0].shape) | |
target_shape[0] = int(target_shape[0] / 2) | |
else: | |
if input_ref_images != None: # Phantom Ref images | |
phantom = True | |
input_ref_images = self.get_vae_latents(input_ref_images, self.device) | |
input_ref_images_neg = torch.zeros_like(input_ref_images) | |
F = frame_num | |
target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1 + (input_ref_images.shape[1] if input_ref_images != None else 0), | |
height // self.vae_stride[1], | |
width // self.vae_stride[2]) | |
seq_len = math.ceil((target_shape[2] * target_shape[3]) / | |
(self.patch_size[1] * self.patch_size[2]) * | |
target_shape[1]) | |
if self._interrupt: | |
return None | |
noise = [ torch.randn( *target_shape, dtype=torch.float32, device=self.device, generator=seed_g) ] | |
# evaluation mode | |
if False: | |
sample_scheduler = FlowMatchScheduler(num_inference_steps=sampling_steps, shift=shift, sigma_min=0, extra_one_step=True) | |
timesteps = torch.tensor([1000, 934, 862, 756, 603, 410, 250, 140, 74, 0])[:sampling_steps].to(self.device) | |
sample_scheduler.timesteps =timesteps | |
elif sample_solver == 'unipc': | |
sample_scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=self.num_train_timesteps, shift=1, use_dynamic_shifting=False) | |
sample_scheduler.set_timesteps( sampling_steps, device=self.device, shift=shift) | |
timesteps = sample_scheduler.timesteps | |
elif sample_solver == 'dpm++': | |
sample_scheduler = FlowDPMSolverMultistepScheduler( | |
num_train_timesteps=self.num_train_timesteps, | |
shift=1, | |
use_dynamic_shifting=False) | |
sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) | |
timesteps, _ = retrieve_timesteps( | |
sample_scheduler, | |
device=self.device, | |
sigmas=sampling_sigmas) | |
else: | |
raise NotImplementedError("Unsupported solver.") | |
# sample videos | |
latents = noise[0] | |
del noise | |
batch_size = 1 | |
if target_camera != None: | |
shape = list(latents.shape[1:]) | |
shape[0] *= 2 | |
freqs = get_rotary_pos_embed(shape, enable_RIFLEx= False) | |
else: | |
freqs = get_rotary_pos_embed(latents.shape[1:], enable_RIFLEx= enable_RIFLEx) | |
kwargs = {'freqs': freqs, 'pipeline': self, 'callback': callback} | |
if target_camera != None: | |
kwargs.update({'cam_emb': cam_emb}) | |
if vace: | |
ref_images_count = len(input_ref_images[0]) if input_ref_images != None and input_ref_images[0] != None else 0 | |
context_scale = context_scale if context_scale != None else [1.0] * len(z) | |
kwargs.update({'vace_context' : z, 'vace_context_scale' : context_scale}) | |
if overlapped_latents != None : | |
overlapped_latents_size = overlapped_latents.shape[1] + 1 | |
# overlapped_latents_size = 3 | |
z_reactive = [ zz[0:16, 0:overlapped_latents_size + ref_images_count].clone() for zz in z] | |
if self.model.enable_cache: | |
x_count = 3 if phantom else 2 | |
self.model.previous_residual = [None] * x_count | |
self.model.compute_teacache_threshold(self.model.cache_start_step, timesteps, self.model.teacache_multiplier) | |
if callback != None: | |
callback(-1, None, True) | |
offload.shared_state["_chipmunk"] = False | |
chipmunk = offload.shared_state.get("_chipmunk", False) | |
if chipmunk: | |
self.model.setup_chipmunk() | |
for i, t in enumerate(tqdm(timesteps)): | |
timestep = [t] | |
if overlapped_latents != None : | |
overlap_noise_factor = overlap_noise / 1000 | |
latent_noise_factor = t / 1000 | |
for zz, zz_r, ll in zip(z, z_reactive, [latents, None]): # extra None for second control net | |
zz[0:16, ref_images_count:overlapped_latents_size + ref_images_count] = zz_r[:, ref_images_count:] * (1.0 - overlap_noise_factor) + torch.randn_like(zz_r[:, ref_images_count:] ) * overlap_noise_factor | |
if ll != None: | |
ll[:, 0:overlapped_latents_size + ref_images_count] = zz_r * (1.0 - latent_noise_factor) + torch.randn_like(zz_r ) * latent_noise_factor | |
if target_camera != None: | |
latent_model_input = torch.cat([latents, source_latents], dim=1) | |
else: | |
latent_model_input = latents | |
kwargs["slg_layers"] = slg_layers if int(slg_start * sampling_steps) <= i < int(slg_end * sampling_steps) else None | |
offload.set_step_no_for_lora(self.model, i) | |
timestep = torch.stack(timestep) | |
kwargs["current_step"] = i | |
kwargs["t"] = timestep | |
if guide_scale == 1: | |
noise_pred = self.model( [latent_model_input], x_id = 0, context = [context], **kwargs)[0] | |
if self._interrupt: | |
return None | |
elif joint_pass: | |
if phantom: | |
pos_it, pos_i, neg = self.model( | |
[ torch.cat([latent_model_input[:,:-input_ref_images.shape[1]], input_ref_images], dim=1) ] * 2 + | |
[ torch.cat([latent_model_input[:,:-input_ref_images_neg.shape[1]], input_ref_images_neg], dim=1)], | |
context = [context, context_null, context_null], **kwargs) | |
else: | |
noise_pred_cond, noise_pred_uncond = self.model( | |
[latent_model_input, latent_model_input], context = [context, context_null], **kwargs) | |
if self._interrupt: | |
return None | |
else: | |
if phantom: | |
pos_it = self.model( | |
[ torch.cat([latent_model_input[:,:-input_ref_images.shape[1]], input_ref_images], dim=1) ], x_id = 0, context = [context], **kwargs | |
)[0] | |
if self._interrupt: | |
return None | |
pos_i = self.model( | |
[ torch.cat([latent_model_input[:,:-input_ref_images.shape[1]], input_ref_images], dim=1) ], x_id = 1, context = [context_null],**kwargs | |
)[0] | |
if self._interrupt: | |
return None | |
neg = self.model( | |
[ torch.cat([latent_model_input[:,:-input_ref_images_neg.shape[1]], input_ref_images_neg], dim=1) ], x_id = 2, context = [context_null], **kwargs | |
)[0] | |
if self._interrupt: | |
return None | |
else: | |
noise_pred_cond = self.model( | |
[latent_model_input], x_id = 0, context = [context], **kwargs)[0] | |
if self._interrupt: | |
return None | |
noise_pred_uncond = self.model( | |
[latent_model_input], x_id = 1, context = [context_null], **kwargs)[0] | |
if self._interrupt: | |
return None | |
# del latent_model_input | |
# CFG Zero *. Thanks to https://github.com/WeichenFan/CFG-Zero-star/ | |
if guide_scale == 1: | |
pass | |
elif phantom: | |
guide_scale_img= 5.0 | |
guide_scale_text= guide_scale #7.5 | |
noise_pred = neg + guide_scale_img * (pos_i - neg) + guide_scale_text * (pos_it - pos_i) | |
else: | |
noise_pred_text = noise_pred_cond | |
if cfg_star_switch: | |
positive_flat = noise_pred_text.view(batch_size, -1) | |
negative_flat = noise_pred_uncond.view(batch_size, -1) | |
alpha = optimized_scale(positive_flat,negative_flat) | |
alpha = alpha.view(batch_size, 1, 1, 1) | |
if (i <= cfg_zero_step): | |
noise_pred = noise_pred_text*0. # it would be faster not to compute noise_pred... | |
else: | |
noise_pred_uncond *= alpha | |
noise_pred = noise_pred_uncond + guide_scale * (noise_pred_text - noise_pred_uncond) | |
noise_pred_uncond, noise_pred_cond, noise_pred_text, pos_it, pos_i, neg = None, None, None, None, None, None | |
scheduler_kwargs = {} if isinstance(sample_scheduler, FlowMatchScheduler) else {"generator": seed_g} | |
temp_x0 = sample_scheduler.step( | |
noise_pred[:, :target_shape[1]].unsqueeze(0), | |
t, | |
latents.unsqueeze(0), | |
# return_dict=False, | |
**scheduler_kwargs)[0] | |
latents = temp_x0.squeeze(0) | |
del temp_x0 | |
if callback is not None: | |
callback(i, latents, False) | |
x0 = [latents] | |
if chipmunk: | |
self.model.release_chipmunk() # need to add it at every exit when in prof | |
if return_latent_slice != None: | |
if overlapped_latents != None: | |
# latents [:, 1:] = self.toto | |
for zz, zz_r, ll in zip(z, z_reactive, [latents]): | |
ll[:, 0:overlapped_latents_size + ref_images_count] = zz_r | |
latent_slice = latents[:, return_latent_slice].clone() | |
if input_frames == None: | |
if phantom: | |
# phantom post processing | |
x0 = [x0_[:,:-input_ref_images.shape[1]] for x0_ in x0] | |
videos = self.vae.decode(x0, VAE_tile_size) | |
else: | |
# vace post processing | |
videos = self.decode_latent(x0, input_ref_images, VAE_tile_size) | |
if return_latent_slice != None: | |
return { "x" : videos[0], "latent_slice" : latent_slice } | |
return videos[0] | |
def adapt_vace_model(self): | |
model = self.model | |
modules_dict= { k: m for k, m in model.named_modules()} | |
for model_layer, vace_layer in model.vace_layers_mapping.items(): | |
module = modules_dict[f"vace_blocks.{vace_layer}"] | |
target = modules_dict[f"blocks.{model_layer}"] | |
setattr(target, "vace", module ) | |
delattr(model, "vace_blocks") |