diff --git a/nodes/acestep/ace_step_guiders.py b/nodes/acestep/ace_step_guiders.py index 7fcc61a..839a181 100644 --- a/nodes/acestep/ace_step_guiders.py +++ b/nodes/acestep/ace_step_guiders.py @@ -140,57 +140,124 @@ class ACEStepExtendGuider(comfy.samplers.CFGGuider): self.left_frames = ACEStepLatentUtils.time_to_frame_index(extend_left_time) self.right_frames = ACEStepLatentUtils.time_to_frame_index(extend_right_time) - # Create extended latent by padding with zeros - self.extended_latent = torch.nn.functional.pad( - source_latent, - (self.left_frames, self.right_frames), - "constant", 0.0 - ) + # Store original source length for post-processing + self.src_latents_length = source_latent.shape[-1] + + # Calculate max inference length (240 seconds = ~2640 frames) + self.max_infer_frame_length = int(240 * 44100 / 512 / 8) + + # Initialize trimming variables (matching community logic) + self.left_trim_length = 0 + self.right_trim_length = 0 + self.to_left_pad_gt_latents = None + self.to_right_pad_gt_latents = None + + # Create extended latent with trimming logic (matching community exactly) + gt_latents = source_latent + frame_length = self.left_frames + gt_latents.shape[-1] + self.right_frames + + # Handle left extension with trimming + if self.left_frames > 0: + extend_gt_latents = torch.nn.functional.pad( + gt_latents, (self.left_frames, 0), "constant", 0 + ) + if frame_length > self.max_infer_frame_length: + self.right_trim_length = frame_length - self.max_infer_frame_length + self.to_right_pad_gt_latents = extend_gt_latents[:, :, :, -self.right_trim_length:] + extend_gt_latents = extend_gt_latents[:, :, :, :self.max_infer_frame_length] + frame_length = self.max_infer_frame_length + gt_latents = extend_gt_latents + + # Handle right extension with trimming + if self.right_frames > 0: + extend_gt_latents = torch.nn.functional.pad( + gt_latents, (0, self.right_frames), "constant", 0 + ) + frame_length = extend_gt_latents.shape[-1] + if frame_length > self.max_infer_frame_length: + self.left_trim_length = frame_length - self.max_infer_frame_length + self.to_left_pad_gt_latents = extend_gt_latents[:, :, :, :self.left_trim_length] + extend_gt_latents = extend_gt_latents[:, :, :, -self.max_infer_frame_length:] + frame_length = self.max_infer_frame_length + gt_latents = extend_gt_latents + + # Store the working latent (this is what we'll use for inference) + self.extended_latent = gt_latents # Create binary mask for extended regions self.extend_mask = torch.zeros_like(self.extended_latent) if self.left_frames > 0: - self.extend_mask[:, :, :, :self.left_frames] = 1.0 + actual_left_frames = min(self.left_frames, self.extended_latent.shape[-1]) + self.extend_mask[:, :, :, :actual_left_frames] = 1.0 if self.right_frames > 0: - self.extend_mask[:, :, :, -self.right_frames:] = 1.0 + actual_right_frames = min(self.right_frames, self.extended_latent.shape[-1]) + self.extend_mask[:, :, :, -actual_right_frames:] = 1.0 # Calculate n_min based on retake_variance (using 1.0 for extend) self.retake_variance = 1.0 def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - """Override sample to implement extend logic with custom ODE""" + """Override sample to implement extend logic exactly matching community implementation""" - # Generate retake_latents (separate noise for extended regions) - retake_latents = torch.randn_like(self.extended_latent) + device = noise.device + dtype = noise.dtype - # Create initial target_latents by concatenating noise pieces + # Generate retake_latents - separate noise for extended regions (matches community) + # Community uses randn_tensor with retake_random_generators, we approximate with randn_like + retake_latents = torch.randn_like(self.extended_latent, device=device, dtype=dtype) + + # Generate target_latents for the FULL extended duration (this is what community does) + # Community generates target_latents for the full frame_length, then slices it + target_latents_full = torch.randn_like(self.extended_latent, device=device, dtype=dtype) + + # Create target_latents by concatenating pieces EXACTLY like community implementation + # Community code lines 768-775: proper slicing with trimming padd_list = [] + + # Left extension: use retake_latents directly if self.left_frames > 0: - left_piece = retake_latents[:, :, :, :self.left_frames] - padd_list.append(left_piece) + actual_left_frames = min(self.left_frames, retake_latents.shape[-1]) + padd_list.append(retake_latents[:, :, :, :actual_left_frames]) - # Middle section: use original noise (trimmed if needed) - original_noise_start = 0 - original_noise_end = noise.shape[-1] - middle_piece = noise[:, :, :, original_noise_start:original_noise_end] - padd_list.append(middle_piece) + # Middle section: use TRIMMED SLICE of target_latents_full (this is the key fix!) + # Community: target_latents[left_trim_length : target_latents.shape[-1] - right_trim_length] + middle_start = self.left_trim_length + middle_end = target_latents_full.shape[-1] - self.right_trim_length + middle_slice = target_latents_full[:, :, :, middle_start:middle_end] + # The middle slice should match the original source latent length + # We need to trim it to match the source latent size + if middle_slice.shape[-1] > self.src_latents_length: + # Trim to source length + middle_slice = middle_slice[:, :, :, :self.src_latents_length] + elif middle_slice.shape[-1] < self.src_latents_length: + # This shouldn't happen in normal cases, but pad if needed + padding_needed = self.src_latents_length - middle_slice.shape[-1] + middle_slice = torch.nn.functional.pad(middle_slice, (0, padding_needed), "constant", 0.0) + + padd_list.append(middle_slice) + + # Right extension: use retake_latents directly if self.right_frames > 0: - right_piece = retake_latents[:, :, :, -self.right_frames:] - padd_list.append(right_piece) + actual_right_frames = min(self.right_frames, retake_latents.shape[-1]) + padd_list.append(retake_latents[:, :, :, -actual_right_frames:]) target_latents = torch.cat(padd_list, dim=-1) + # Ensure target_latents matches extended_latent shape + assert target_latents.shape[-1] == self.extended_latent.shape[-1], \ + f"Shape mismatch: {target_latents.shape=} vs {self.extended_latent.shape=}" + # Calculate n_min total_steps = len(sigmas) - 1 n_min = int(total_steps * (1 - self.retake_variance)) - # Core variables - x0 = self.extended_latent # ground truth latents - z0 = target_latents # initial noise - zt_edit = x0.clone() # will be updated - is_repaint = True # extend is treated as repaint - repaint_mask = self.extend_mask + # Core variables - exactly matching community variable names and assignments + x0 = self.extended_latent # gt_latents in community + z0 = target_latents # target_latents in community (becomes z0) + zt_edit = x0.clone() # zt_edit = x0.clone() in community + is_repaint = True # extend is treated as repaint in community + repaint_mask = self.extend_mask # repaint_mask in community # Convert sigmas to timesteps timesteps = (sigmas * 1000).long() @@ -199,20 +266,21 @@ class ACEStepExtendGuider(comfy.samplers.CFGGuider): original_dtype = target_latents.dtype def extend_callback(step, x0_unused, x, total_steps): - """Implement extend logic with custom ODE""" + """Implement extend logic exactly matching community ODE""" nonlocal target_latents, zt_edit device = x.device i = step t = timesteps[step] if step < len(timesteps) else timesteps[-1] - # Repaint initialization + # Community repaint initialization logic if is_repaint: if i < n_min: - # Let normal scheduler work for steps before n_min + # Let normal scheduler work for steps before n_min (community: continue) target_latents = x.clone() return elif i == n_min: + # Community initialization at n_min t_i = t.float() / 1000.0 zt_src = (1 - t_i) * x0.to(device) + t_i * z0.to(device) target_latents = zt_edit.to(device) + zt_src - x0.to(device) @@ -221,33 +289,30 @@ class ACEStepExtendGuider(comfy.samplers.CFGGuider): x[:] = target_latents return - # For steps after n_min: BYPASS SCHEDULER and do custom ODE + # Community ODE step for i >= n_min if is_repaint and i >= n_min and i > 0: - # Get timesteps for ODE integration + # Get timesteps for ODE integration (matches community exactly) t_i = t.float() / 1000.0 if i + 1 < len(timesteps): t_im1 = timesteps[i + 1].float() / 1000.0 else: t_im1 = torch.zeros_like(t_i).to(t_i.device) - # Dtype conversion for precision + # Community ODE step: target_latents = target_latents + (t_im1 - t_i) * noise_pred + # But we're in callback after scheduler step, so x already contains the result + # We use x as our prev_sample (matches community's prev_sample) + dtype = x.dtype # Use x's dtype (matches community) target_latents = target_latents.to(torch.float32) - - # Apply custom ODE step - # Note: The scheduler has already done its step, so x contains the "prev_sample" - # We'll use this as our ODE result for the generated regions prev_sample = x.to(torch.float32) - - # Convert back to original dtype - prev_sample = prev_sample.to(original_dtype) + prev_sample = prev_sample.to(dtype) target_latents = prev_sample - # preserve source regions, keep generated in extended regions + # Community masking step: preserve source regions, keep generated in extended regions zt_src = (1 - t_im1) * x0.to(device) + t_im1 * z0.to(device) - extend_mask = repaint_mask.to(device) - target_latents = torch.where(extend_mask == 1.0, target_latents, zt_src) + repaint_mask_device = repaint_mask.to(device) + target_latents = torch.where(repaint_mask_device == 1.0, target_latents, zt_src) - # Replace x with our custom ODE result + # Replace x with our result x[:] = target_latents # Call original callback if provided @@ -257,6 +322,14 @@ class ACEStepExtendGuider(comfy.samplers.CFGGuider): # Use extended latent and target_latents as initial noise result = super().sample(target_latents, self.extended_latent, sampler, sigmas, denoise_mask, extend_callback, disable_pbar, seed) + # POST-PROCESSING: Apply community's final concatenation logic (lines 1140-1147) + # This is the critical missing piece that makes extend sound like continuation! + if self.to_right_pad_gt_latents is not None: + result = torch.cat([result, self.to_right_pad_gt_latents.to(result.device)], dim=-1) + if self.to_left_pad_gt_latents is not None: + # Fix the bug in community code: should be [to_left_pad_gt_latents, result] not [to_right_pad_gt_latents, result] + result = torch.cat([self.to_left_pad_gt_latents.to(result.device), result], dim=-1) + return result class ACEStepHybridGuider(comfy.samplers.CFGGuider): @@ -281,27 +354,70 @@ class ACEStepHybridGuider(comfy.samplers.CFGGuider): self.has_repaint = repaint_start_time is not None and repaint_end_time is not None self.has_extend = extend_left_time > 0 or extend_right_time > 0 - # Setup extend if enabled + # Setup extend if enabled (using improved logic from ACEStepExtendGuider) if self.has_extend: self.left_frames = ACEStepLatentUtils.time_to_frame_index(extend_left_time) self.right_frames = ACEStepLatentUtils.time_to_frame_index(extend_right_time) - # Create extended latent by padding with zeros - self.extended_latent = torch.nn.functional.pad( - source_latent, - (self.left_frames, self.right_frames), - "constant", 0.0 - ) + # Store original source length for post-processing + self.src_latents_length = source_latent.shape[-1] + + # Calculate max inference length (240 seconds = ~2640 frames) + self.max_infer_frame_length = int(240 * 44100 / 512 / 8) + + # Initialize trimming variables (matching community logic) + self.left_trim_length = 0 + self.right_trim_length = 0 + self.to_left_pad_gt_latents = None + self.to_right_pad_gt_latents = None + + # Create extended latent with trimming logic (matching community exactly) + gt_latents = source_latent + frame_length = self.left_frames + gt_latents.shape[-1] + self.right_frames + + # Handle left extension with trimming + if self.left_frames > 0: + extend_gt_latents = torch.nn.functional.pad( + gt_latents, (self.left_frames, 0), "constant", 0 + ) + if frame_length > self.max_infer_frame_length: + self.right_trim_length = frame_length - self.max_infer_frame_length + self.to_right_pad_gt_latents = extend_gt_latents[:, :, :, -self.right_trim_length:] + extend_gt_latents = extend_gt_latents[:, :, :, :self.max_infer_frame_length] + frame_length = self.max_infer_frame_length + gt_latents = extend_gt_latents + + # Handle right extension with trimming + if self.right_frames > 0: + extend_gt_latents = torch.nn.functional.pad( + gt_latents, (0, self.right_frames), "constant", 0 + ) + frame_length = extend_gt_latents.shape[-1] + if frame_length > self.max_infer_frame_length: + self.left_trim_length = frame_length - self.max_infer_frame_length + self.to_left_pad_gt_latents = extend_gt_latents[:, :, :, :self.left_trim_length] + extend_gt_latents = extend_gt_latents[:, :, :, -self.max_infer_frame_length:] + frame_length = self.max_infer_frame_length + gt_latents = extend_gt_latents + + # Store the working latent (this is what we'll use for inference) + self.extended_latent = gt_latents # Create binary mask for extended regions self.extend_mask = torch.zeros_like(self.extended_latent) if self.left_frames > 0: - self.extend_mask[:, :, :, :self.left_frames] = 1.0 + actual_left_frames = min(self.left_frames, self.extended_latent.shape[-1]) + self.extend_mask[:, :, :, :actual_left_frames] = 1.0 if self.right_frames > 0: - self.extend_mask[:, :, :, -self.right_frames:] = 1.0 + actual_right_frames = min(self.right_frames, self.extended_latent.shape[-1]) + self.extend_mask[:, :, :, -actual_right_frames:] = 1.0 else: self.extended_latent = source_latent self.extend_mask = None + self.left_trim_length = 0 + self.right_trim_length = 0 + self.to_left_pad_gt_latents = None + self.to_right_pad_gt_latents = None # Setup repaint if enabled if self.has_repaint: @@ -342,11 +458,49 @@ class ACEStepHybridGuider(comfy.samplers.CFGGuider): if hasattr(noise, 'seed'): working_noise = comfy.sample.prepare_noise(working_latent, noise.seed) else: - working_noise = torch.nn.functional.pad( - noise, - (self.left_frames, self.right_frames), - "constant", 0.0 - ) + # Use improved noise generation strategy for better extend blending + device = noise.device + dtype = noise.dtype + + # Generate additional noise for extended regions + retake_noise = torch.randn_like(self.extended_latent, device=device, dtype=dtype) + + # Generate target_latents for the FULL extended duration (matches community) + target_latents_full = torch.randn_like(self.extended_latent, device=device, dtype=dtype) + + # Create noise by concatenating pieces EXACTLY like community implementation + # Community does NOT use cosine/sine mixing for extend - only direct concatenation + padd_list = [] + + # Left extension: use retake_noise directly + if self.left_frames > 0: + actual_left_frames = min(self.left_frames, retake_noise.shape[-1]) + padd_list.append(retake_noise[:, :, :, :actual_left_frames]) + + # Middle section: use TRIMMED SLICE of target_latents_full (key fix!) + middle_start = self.left_trim_length + middle_end = target_latents_full.shape[-1] - self.right_trim_length + middle_slice = target_latents_full[:, :, :, middle_start:middle_end] + + # Trim to source length + if middle_slice.shape[-1] > self.src_latents_length: + middle_slice = middle_slice[:, :, :, :self.src_latents_length] + elif middle_slice.shape[-1] < self.src_latents_length: + padding_needed = self.src_latents_length - middle_slice.shape[-1] + middle_slice = torch.nn.functional.pad(middle_slice, (0, padding_needed), "constant", 0.0) + + padd_list.append(middle_slice) + + # Right extension: use retake_noise directly + if self.right_frames > 0: + actual_right_frames = min(self.right_frames, retake_noise.shape[-1]) + padd_list.append(retake_noise[:, :, :, -actual_right_frames:]) + + working_noise = torch.cat(padd_list, dim=-1) + + # Ensure working_noise matches working_latent shape + assert working_noise.shape[-1] == working_latent.shape[-1], \ + f"Shape mismatch: {working_noise.shape=} vs {working_latent.shape=}" else: working_noise = noise @@ -427,4 +581,14 @@ class ACEStepHybridGuider(comfy.samplers.CFGGuider): return callback(step, x0_unused, x, total_steps) # Run sampling with working latent and hybrid callback - return super().sample(working_noise, working_latent, sampler, sigmas, denoise_mask, hybrid_callback, disable_pbar, seed) \ No newline at end of file + result = super().sample(working_noise, working_latent, sampler, sigmas, denoise_mask, hybrid_callback, disable_pbar, seed) + + # POST-PROCESSING: Apply community's final concatenation logic (if extend is enabled) + if self.has_extend: + if self.to_right_pad_gt_latents is not None: + result = torch.cat([result, self.to_right_pad_gt_latents.to(result.device)], dim=-1) + if self.to_left_pad_gt_latents is not None: + # Fix the bug in community code: should be [to_left_pad_gt_latents, result] not [to_right_pad_gt_latents, result] + result = torch.cat([self.to_left_pad_gt_latents.to(result.device), result], dim=-1) + + return result \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 8a977e2..2434d64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_ryanonyheinside" description = "Custom nodes introducing everything-reactivity and timelines to ComfyUI" -version = "2.0.10" +version = "2.0.11" license = {file = "LICENSE"} dependencies = ["pygame","opencv-python>=4.10.0","scipy","torchaudio", "pillow","librosa>=0.10.2","pymunk==6.8.1","matplotlib","openunmix","mido","scikit-image"]