@@ -316,8 +316,16 @@ def prepare_latents(
316316
317317 ref_latent = self .vae .encode (ref_tensor .to (self .vae .dtype )).latent_dist .mode ()
318318 ref_latent = ref_latent .to (dtype )
319- latents_mean = torch .tensor (self .vae .config .latents_mean ).view (1 , - 1 , 1 , 1 , 1 ).to (ref_latent .device , ref_latent .dtype )
320- latents_std = torch .tensor (self .vae .config .latents_std ).view (1 , - 1 , 1 , 1 , 1 ).to (ref_latent .device , ref_latent .dtype )
319+ latents_mean = (
320+ torch .tensor (self .vae .config .latents_mean )
321+ .view (1 , - 1 , 1 , 1 , 1 )
322+ .to (ref_latent .device , ref_latent .dtype )
323+ )
324+ latents_std = (
325+ torch .tensor (self .vae .config .latents_std )
326+ .view (1 , - 1 , 1 , 1 , 1 )
327+ .to (ref_latent .device , ref_latent .dtype )
328+ )
321329 ref_latent = (ref_latent - latents_mean ) / latents_std
322330 ref_latent = ref_latent .squeeze (0 ) # [C, 1, H', W']
323331 sample_items .append (ref_latent )
@@ -674,8 +682,16 @@ def __call__(
674682 1 , c_lat , l_t * pt , l_h * ph , l_w * pw
675683 )
676684
677- latents_mean = torch .tensor (self .vae .config .latents_mean ).view (1 , - 1 , 1 , 1 , 1 ).to (video_latent .device , video_latent .dtype )
678- latents_std = torch .tensor (self .vae .config .latents_std ).view (1 , - 1 , 1 , 1 , 1 ).to (video_latent .device , video_latent .dtype )
685+ latents_mean = (
686+ torch .tensor (self .vae .config .latents_mean )
687+ .view (1 , - 1 , 1 , 1 , 1 )
688+ .to (video_latent .device , video_latent .dtype )
689+ )
690+ latents_std = (
691+ torch .tensor (self .vae .config .latents_std )
692+ .view (1 , - 1 , 1 , 1 , 1 )
693+ .to (video_latent .device , video_latent .dtype )
694+ )
679695 video_latent = video_latent * latents_std + latents_mean
680696
681697 sample_image = self .vae .decode (video_latent .to (self .vae .dtype ), return_dict = False )[0 ]
0 commit comments