diff --git a/latent_preview.py b/latent_preview.py index fb94e54..8d51f8e 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -84,12 +84,12 @@ def get_previewer(device, latent_format): taew_sd = comfy.utils.load_torch_file(taehv_path) taesd = TAEHV(taew_sd).to(device) previewer = TAESDPreviewerImpl(taesd) - previewer = WrappedPreviewer(previewer) + previewer = WrappedPreviewer(previewer, rate=16) if previewer is None: if latent_format.latent_rgb_factors is not None: previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors, latent_format.latent_rgb_factors_bias) - previewer = WrappedPreviewer(previewer) + previewer = WrappedPreviewer(previewer, rate=4) return previewer def prepare_callback(model, steps, x0_output_dict=None): @@ -182,7 +182,10 @@ class WrappedPreviewer(LatentPreviewer): #NOTE: send sync already uses call_soon_threadsafe serv.send_sync(server.BinaryEventTypes.PREVIEW_IMAGE, message.getvalue(), serv.client_id) - ind = (ind + 1) % ((leng-1) * 4 - 1) + if self.rate == 16: + ind = (ind + 1) % ((leng-1) * 4 - 1) + else: + ind = (ind + 1) % leng def decode_latent_to_preview(self, x0): if hasattr(self, 'taesd'): x0 = x0.unsqueeze(0) diff --git a/nodes.py b/nodes.py index 268ff46..e955e22 100644 --- a/nodes.py +++ b/nodes.py @@ -1694,6 +1694,7 @@ class WanVideoSampler: ) mask = masks[idx] mask = mask.to(latent) + print(mask.shape, image_latent.shape, latent.shape) latent = image_latent * mask + latent * (1-mask) # end diff diff @@ -2105,33 +2106,21 @@ class WanVideoEncode: vae.to(offload_device) mm.soft_empty_cache() - print("encoded latents shape",latents.shape) + log.info(f"encoded latents shape {latents.shape}") if mask is not None: #B, H, W - B, H, W = mask.shape - target_frames = latents.shape[2] target_h, target_w = latents.shape[3:] - - # Temporal: pad/truncate - if B > target_frames: - mask = mask[:target_frames] - elif B < target_frames: - padding = torch.zeros((target_frames - B, H, W), device=mask.device) - mask = torch.cat([mask, padding], dim=0) - - # Spatial: resize each frame + mask = torch.nn.functional.interpolate( - mask.unsqueeze(1), # Add channel dim for interpolate - size=(target_h, target_w), - mode='bilinear' - ).squeeze(1) # Remove channel dim + mask.unsqueeze(0).unsqueeze(0), # Add batch and channel dims [1,1,T,H,W] + size=(latents.shape[2], target_h, target_w), + mode='trilinear', + align_corners=False + ).squeeze(0) # Remove batch dim, keep channel dim # Add batch & channel dims for final output - mask = mask.unsqueeze(0).unsqueeze(0) - mask = mask.repeat(1, latents.shape[1], 1, 1, 1) - print("mask shape",mask.shape) - - + mask = mask.unsqueeze(0).repeat(1, latents.shape[1], 1, 1, 1) + return ({"samples": latents, "mask": mask},) class WanVideoLatentPreview: