Spaces:
Running
Running
| """Padding helpers for handler batch preparation.""" | |
| import torch | |
| from loguru import logger | |
| class PaddingMixin: | |
| """Mixin containing repaint/lego padding helpers. | |
| Depends on host members: | |
| - Method: ``create_target_wavs`` (provided by ``TaskUtilsMixin`` in this decomposition). | |
| """ | |
| def prepare_padding_info( | |
| self, | |
| actual_batch_size, | |
| processed_src_audio, | |
| audio_duration, | |
| repainting_start, | |
| repainting_end, | |
| is_repaint_task, | |
| is_lego_task, | |
| is_cover_task, | |
| can_use_repainting, | |
| ): | |
| """Prepare padded target wavs and repaint coordinates for each batch item.""" | |
| try: | |
| target_wavs_batch = [] | |
| # Store padding info for each batch item to adjust repainting coordinates | |
| padding_info_batch = [] | |
| for i in range(actual_batch_size): | |
| if processed_src_audio is not None: | |
| if is_cover_task: | |
| # Cover task: Use src_audio directly without padding | |
| batch_target_wavs = processed_src_audio | |
| padding_info_batch.append({"left_padding_duration": 0.0, "right_padding_duration": 0.0}) | |
| elif is_repaint_task or is_lego_task: | |
| # Repaint/lego task: May need padding for outpainting | |
| src_audio_duration = processed_src_audio.shape[-1] / 48000.0 | |
| # Determine actual end time | |
| if repainting_end is None or repainting_end < 0: | |
| actual_end = src_audio_duration | |
| else: | |
| actual_end = repainting_end | |
| left_padding_duration = max(0, -repainting_start) if repainting_start is not None else 0 | |
| right_padding_duration = max(0, actual_end - src_audio_duration) | |
| # Create padded audio | |
| left_padding_frames = int(left_padding_duration * 48000) | |
| right_padding_frames = int(right_padding_duration * 48000) | |
| if left_padding_frames > 0 or right_padding_frames > 0: | |
| # Pad the src audio | |
| batch_target_wavs = torch.nn.functional.pad( | |
| processed_src_audio, (left_padding_frames, right_padding_frames), "constant", 0 | |
| ) | |
| else: | |
| batch_target_wavs = processed_src_audio | |
| # Store padding info for coordinate adjustment | |
| padding_info_batch.append( | |
| { | |
| "left_padding_duration": left_padding_duration, | |
| "right_padding_duration": right_padding_duration, | |
| } | |
| ) | |
| else: | |
| # Other tasks: Use src_audio directly without padding | |
| batch_target_wavs = processed_src_audio | |
| padding_info_batch.append({"left_padding_duration": 0.0, "right_padding_duration": 0.0}) | |
| else: | |
| padding_info_batch.append({"left_padding_duration": 0.0, "right_padding_duration": 0.0}) | |
| if audio_duration is not None and float(audio_duration) > 0: | |
| batch_target_wavs = self.create_target_wavs(float(audio_duration)) | |
| else: | |
| import random | |
| random_duration = random.uniform(10.0, 120.0) | |
| batch_target_wavs = self.create_target_wavs(random_duration) | |
| target_wavs_batch.append(batch_target_wavs) | |
| # Stack target_wavs into batch tensor | |
| # Ensure all tensors have the same shape by padding to max length | |
| max_frames = max(wav.shape[-1] for wav in target_wavs_batch) | |
| padded_target_wavs = [] | |
| for wav in target_wavs_batch: | |
| if wav.shape[-1] < max_frames: | |
| pad_frames = max_frames - wav.shape[-1] | |
| padded_wav = torch.nn.functional.pad(wav, (0, pad_frames), "constant", 0) | |
| padded_target_wavs.append(padded_wav) | |
| else: | |
| padded_target_wavs.append(wav) | |
| target_wavs_tensor = torch.stack(padded_target_wavs, dim=0) # [batch_size, 2, frames] | |
| if can_use_repainting: | |
| # Repaint task: Set repainting parameters | |
| if repainting_start is None: | |
| repainting_start_batch = None | |
| elif isinstance(repainting_start, (int, float)): | |
| if processed_src_audio is not None: | |
| adjusted_start = repainting_start + padding_info_batch[0]["left_padding_duration"] | |
| repainting_start_batch = [adjusted_start] * actual_batch_size | |
| else: | |
| repainting_start_batch = [repainting_start] * actual_batch_size | |
| else: | |
| # List input - adjust each item | |
| repainting_start_batch = [] | |
| for i in range(actual_batch_size): | |
| if processed_src_audio is not None: | |
| adjusted_start = repainting_start[i] + padding_info_batch[i]["left_padding_duration"] | |
| repainting_start_batch.append(adjusted_start) | |
| else: | |
| repainting_start_batch.append(repainting_start[i]) | |
| # Handle repainting_end - use src audio duration if not specified or negative | |
| if processed_src_audio is not None: | |
| # If src audio is provided, use its duration as default end | |
| src_audio_duration = processed_src_audio.shape[-1] / 48000.0 | |
| if repainting_end is None or repainting_end < 0: | |
| # Use src audio duration (before padding), then adjust for padding | |
| adjusted_end = src_audio_duration + padding_info_batch[0]["left_padding_duration"] | |
| repainting_end_batch = [adjusted_end] * actual_batch_size | |
| else: | |
| # Adjust repainting_end to be relative to padded audio | |
| adjusted_end = repainting_end + padding_info_batch[0]["left_padding_duration"] | |
| repainting_end_batch = [adjusted_end] * actual_batch_size | |
| else: | |
| # No src audio - repainting doesn't make sense without it | |
| if repainting_end is None or repainting_end < 0: | |
| repainting_end_batch = None | |
| elif isinstance(repainting_end, (int, float)): | |
| repainting_end_batch = [repainting_end] * actual_batch_size | |
| else: | |
| # List input - adjust each item | |
| repainting_end_batch = [] | |
| for i in range(actual_batch_size): | |
| repainting_end_batch.append(repainting_end[i]) | |
| else: | |
| # All other tasks (cover, text2music, extract, complete): No repainting | |
| # Only repaint and lego tasks should have repainting parameters | |
| repainting_start_batch = None | |
| repainting_end_batch = None | |
| return repainting_start_batch, repainting_end_batch, target_wavs_tensor | |
| except (TypeError, ValueError, RuntimeError, IndexError): | |
| logger.exception("[prepare_padding_info] Error preparing padding information") | |
| fallback = torch.stack([self.create_target_wavs(30.0) for _ in range(actual_batch_size)], dim=0) | |
| return None, None, fallback | |