Function process_data_batch

pipeline/code/06_superres_apollo.py:88–205

(audio_items, model, output_dir)

        return None, None, None, None, None, None, None, None, None

def process_data_batch(audio_items, model, output_dir):
    results = []
    
    all_chunks = []
    audio_infos = []
    
    valid_items = []
    for item in audio_items:
        audio_path = item["audio_path"]
        if not os.path.exists(audio_path):
            logger.warning(f"Audio file not found: {audio_path}")
            continue
        valid_items.append(item)
    
    if not valid_items:
        return results
    
    max_workers = min(len(valid_items), 8)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_item = {executor.submit(prepare_audio_chunks, item["audio_path"]): item 
                         for item in valid_items}
        
        for future in concurrent.futures.as_completed(future_to_item):
            item = future_to_item[future]
            try:
                chunks, positions, shape, orig_len, sr, C, step, fade_size, border = future.result()
                if chunks is None:
                    continue
                    
                audio_infos.append({
                    'item': item,
                    'chunks_start_idx': len(all_chunks),
                    'chunks_count': len(chunks),
                    'positions': positions,
                    'shape': shape,
                    'original_length': orig_len,
                    'samplerate': sr,
                    'C': C,
                    'step': step,
                    'fade_size': fade_size,
                    'border': border
                })
                
                all_chunks.extend(chunks)
            except Exception as e:
                logger.error(f"Error processing audio file {item['audio_path']}: {e}")
                continue
    
    if not all_chunks:
        return results
    
    try:
        batch_chunks = torch.stack(all_chunks).to(device)
        with torch.no_grad():
            batch_outputs = model(batch_chunks).squeeze(1).cpu()
    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            logger.error(f"CUDA out of memory: {e}")