Added RIFLEx support
This commit is contained in:
DeepBeepMeep
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user