27 lines
676 B
Python
27 lines
676 B
Python
import torch
|
|
|
|
|
|
def get_latent_dimensions(num_frames, width, height):
|
|
spatial_downsample = 8
|
|
temporal_downsample = 6
|
|
in_channels = 12
|
|
B = 1
|
|
C = in_channels
|
|
T = (num_frames - 1) // temporal_downsample + 1
|
|
H = height // spatial_downsample
|
|
W = width // spatial_downsample
|
|
|
|
return (B, C, T, H, W)
|
|
|
|
|
|
def add_latent_noise(model, latent_shape, sigma_schedule, samples, generator):
|
|
z = torch.randn(
|
|
latent_shape,
|
|
device=model.device,
|
|
generator=generator,
|
|
dtype=torch.float32,
|
|
)
|
|
if samples is not None:
|
|
z = z * sigma_schedule[0] + (1 -sigma_schedule[0]) * samples.to(model.device)
|
|
return z
|
|
|