comi-logic/Wan2GP
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added RIFLEx support

Browse Source
This commit is contained in:
DeepBeepMeep
2025-03-02 16:46:52 +01:00
parent 4203fd3732
commit 3731ab70e1
6 changed files with 317 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

183
wan/image2video.py
View File

@@ -143,7 +143,9 @@ class WanI2V:
n_prompt="",
seed=-1,
offload_model=True,
callback = None
callback = None,
enable_RIFLEx = False
):
r"""
Generates video frames from input image and text prompt using diffusion process.
@@ -262,104 +264,107 @@ class WanI2V:
no_sync = getattr(self.model, 'no_sync', noop_no_sync)
# evaluation mode
with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
if 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.")
if 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
latent = noise
# sample videos
latent = noise
arg_c = {
'context': [context[0]],
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'pipeline' : self
}
freqs = self.model.get_rope_freqs(nb_latent_frames = int((frame_num - 1)/4 + 1), RIFLEx_k = 4 if enable_RIFLEx else None )
arg_null = {
'context': context_null,
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'pipeline' : self
}
arg_c = {
'context': [context[0]],
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'freqs' : freqs,
'pipeline' : self
}
arg_null = {
'context': context_null,
'clip_fea': clip_context,
'seq_len': max_seq_len,
'y': [y],
'freqs' : freqs,
'pipeline' : self
}
if offload_model:
torch.cuda.empty_cache()
# self.model.to(self.device)
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = [latent.to(self.device)]
timestep = [t]
timestep = torch.stack(timestep).to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
if offload_model:
torch.cuda.empty_cache()
# self.model.to(self.device)
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = [latent.to(self.device)]
timestep = [t]
timestep = torch.stack(timestep).to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
if offload_model:
torch.cuda.empty_cache()
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
del latent_model_input
if offload_model:
torch.cuda.empty_cache()
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
latent = latent.to(
torch.device('cpu') if offload_model else self.device)
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
t,
latent.unsqueeze(0),
return_dict=False,
generator=seed_g)[0]
latent = temp_x0.squeeze(0)
del temp_x0
del timestep
if callback is not None:
callback(i, latent)
x0 = [latent.to(self.device)]
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
del latent_model_input
if offload_model:
self.model.cpu()
torch.cuda.empty_cache()
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
if self.rank == 0:
videos = self.vae.decode(x0)
latent = latent.to(
torch.device('cpu') if offload_model else self.device)
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
t,
latent.unsqueeze(0),
return_dict=False,
generator=seed_g)[0]
latent = temp_x0.squeeze(0)
del temp_x0
del timestep
if callback is not None:
callback(i, latent)
x0 = [latent.to(self.device)]
if offload_model:
self.model.cpu()
torch.cuda.empty_cache()
if self.rank == 0:
videos = self.vae.decode(x0)
del noise, latent
del sample_scheduler

84
wan/modules/model.py
View File

@@ -6,6 +6,8 @@ import torch.cuda.amp as amp
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
import numpy as np
from typing import Union,Optional
from .attention import pay_attention
@@ -25,7 +27,49 @@ def sinusoidal_embedding_1d(dim, position):
return x
# @amp.autocast(enabled=False)
def identify_k( b: float, d: int, N: int):
"""
This function identifies the index of the intrinsic frequency component in a RoPE-based pre-trained diffusion transformer.
Args:
b (`float`): The base frequency for RoPE.
d (`int`): Dimension of the frequency tensor
N (`int`): the first observed repetition frame in latent space
Returns:
k (`int`): the index of intrinsic frequency component
N_k (`int`): the period of intrinsic frequency component in latent space
Example:
In HunyuanVideo, b=256 and d=16, the repetition occurs approximately 8s (N=48 in latent space).
k, N_k = identify_k(b=256, d=16, N=48)
In this case, the intrinsic frequency index k is 4, and the period N_k is 50.
"""
# Compute the period of each frequency in RoPE according to Eq.(4)
periods = []
for j in range(1, d // 2 + 1):
theta_j = 1.0 / (b ** (2 * (j - 1) / d))
N_j = round(2 * torch.pi / theta_j)
periods.append(N_j)
# Identify the intrinsic frequency whose period is closed to N（see Eq.(7)）
diffs = [abs(N_j - N) for N_j in periods]
k = diffs.index(min(diffs)) + 1
N_k = periods[k-1]
return k, N_k
def rope_params_riflex(max_seq_len, dim, theta=10000, L_test=30, k=6):
assert dim % 2 == 0
exponents = torch.arange(0, dim, 2, dtype=torch.float64).div(dim)
inv_theta_pow = 1.0 / torch.pow(theta, exponents)
inv_theta_pow[k-1] = 0.9 * 2 * torch.pi / L_test
freqs = torch.outer(torch.arange(max_seq_len), inv_theta_pow)
freqs = torch.polar(torch.ones_like(freqs), freqs)
return freqs
def rope_params(max_seq_len, dim, theta=10000):
assert dim % 2 == 0
freqs = torch.outer(
@@ -588,14 +632,6 @@ class WanModel(ModelMixin, ConfigMixin):
self.head = Head(dim, out_dim, patch_size, eps)
# buffers (don't use register_buffer otherwise dtype will be changed in to())
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
d = dim // num_heads
self.freqs = torch.cat([
rope_params(1024, d - 4 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
rope_params(1024, 2 * (d // 6))
],
dim=1)
if model_type == 'i2v':
self.img_emb = MLPProj(1280, dim)
@@ -603,6 +639,29 @@ class WanModel(ModelMixin, ConfigMixin):
# initialize weights
self.init_weights()
# self.freqs = torch.cat([
# rope_params(1024, d - 4 * (d // 6)), #44
# rope_params(1024, 2 * (d // 6)), #42
# rope_params(1024, 2 * (d // 6)) #42
# ],dim=1)
def get_rope_freqs(self, nb_latent_frames, RIFLEx_k = None):
dim = self.dim
num_heads = self.num_heads
d = dim // num_heads
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
freqs = torch.cat([
rope_params_riflex(1024, dim= d - 4 * (d // 6), L_test=nb_latent_frames, k = RIFLEx_k ), #44
rope_params(1024, 2 * (d // 6)), #42
rope_params(1024, 2 * (d // 6)) #42
],dim=1)
return freqs
def forward(
self,
x,
@@ -611,6 +670,7 @@ class WanModel(ModelMixin, ConfigMixin):
seq_len,
clip_fea=None,
y=None,
freqs = None,
pipeline = None,
):
r"""
@@ -638,8 +698,8 @@ class WanModel(ModelMixin, ConfigMixin):
assert clip_fea is not None and y is not None
# params
device = self.patch_embedding.weight.device
if self.freqs.device != device:
self.freqs = self.freqs.to(device)
if freqs.device != device:
freqs = freqs.to(device)
if y is not None:
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
@@ -683,7 +743,7 @@ class WanModel(ModelMixin, ConfigMixin):
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
freqs=freqs,
context=context,
context_lens=context_lens)

132
wan/text2video.py
View File

@@ -128,7 +128,8 @@ class WanT2V:
n_prompt="",
seed=-1,
offload_model=True,
callback = None
callback = None,
enable_RIFLEx = None
):
r"""
Generates video frames from text prompt using diffusion process.
@@ -209,77 +210,84 @@ class WanT2V:
no_sync = getattr(self.model, 'no_sync', noop_no_sync)
# evaluation mode
with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
if 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.")
if 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
# sample videos
latents = noise
arg_c = {'context': context, 'seq_len': seq_len, 'pipeline': self}
arg_null = {'context': context_null, 'seq_len': seq_len, 'pipeline': self}
# from .modules.model import identify_k
# for nf in range(20, 50):
# k, N_k = identify_k(10000, 44, 26)
# print(f"value nb latent frames={nf}, k={k}, n_k={N_k}")
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = latents
timestep = [t]
freqs = self.model.get_rope_freqs(nb_latent_frames = int((frame_num - 1)/4 + 1), RIFLEx_k = 4 if enable_RIFLEx else None )
timestep = torch.stack(timestep)
arg_c = {'context': context, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
arg_null = {'context': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
# self.model.to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
del latent_model_input
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
if callback != None:
callback(-1, None)
self._interrupt = False
for i, t in enumerate(tqdm(timesteps)):
latent_model_input = latents
timestep = [t]
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
t,
latents[0].unsqueeze(0),
return_dict=False,
generator=seed_g)[0]
latents = [temp_x0.squeeze(0)]
del temp_x0
timestep = torch.stack(timestep)
if callback is not None:
callback(i, latents)
# self.model.to(self.device)
noise_pred_cond = self.model(
latent_model_input, t=timestep, **arg_c)[0]
if self._interrupt:
return None
noise_pred_uncond = self.model(
latent_model_input, t=timestep, **arg_null)[0]
if self._interrupt:
return None
x0 = latents
if offload_model:
self.model.cpu()
torch.cuda.empty_cache()
if self.rank == 0:
videos = self.vae.decode(x0)
del latent_model_input
noise_pred = noise_pred_uncond + guide_scale * (
noise_pred_cond - noise_pred_uncond)
del noise_pred_uncond
temp_x0 = sample_scheduler.step(
noise_pred.unsqueeze(0),
t,
latents[0].unsqueeze(0),
return_dict=False,
generator=seed_g)[0]
latents = [temp_x0.squeeze(0)]
del temp_x0
if callback is not None:
callback(i, latents)
x0 = latents
if offload_model:
self.model.cpu()
torch.cuda.empty_cache()
if self.rank == 0:
videos = self.vae.decode(x0)
del noise, latents