Function encode

preprocessing/dataset_tools.py:451–529

Encode pixel data to VAE latents. Parallelization and Memory Management: Use --gpus to control the number of GPUs used for parallel encoding (default: 8). Use --batch-size to control how many images each GPU processes in each batch (default: 100). Images are distributed acr

(
    model_url: str,
    source: str,
    dest: str,
    max_images: Optional[int],
    gpus: int,
    batch_size: int,
)

@click.option('--batch-size', help='Number of images per GPU in each batch', metavar='INT', type=int, default=100, show_default=True)

def encode(
    model_url: str,
    source: str,
    dest: str,
    max_images: Optional[int],
    gpus: int,
    batch_size: int,
):
    """Encode pixel data to VAE latents.
    
    Parallelization and Memory Management:
    
    Use --gpus to control the number of GPUs used for parallel encoding
    (default: 8). Use --batch-size to control how many images each GPU
    processes in each batch (default: 100). Images are distributed across
    GPUs in round-robin fashion and processed in batches to avoid loading
    all images into memory at once.
    
    Example:
    \b
    python dataset_tool.py encode --source=datasets/img64.zip \\
        --dest=datasets/img64_encoded.zip --gpus=8 --batch-size=50
    """
    PIL.Image.init()
    if dest == '':
        raise click.ClickException('--dest output filename or directory must not be an empty string')

    num_files, input_iter = open_dataset(source, max_images=max_images)
    archive_root_dir, save_bytes, close_dest = open_dest(dest)
    
    # Process images in batches across GPUs to avoid loading everything into memory
    labels = []
    
    print(f"Processing {num_files} images in batches of {batch_size * gpus} across {gpus} GPUs...")
    
    with mp.Pool(gpus) as pool:
        batch = []
        gpu_batches = [[] for _ in range(gpus)]
        
        for idx, image in tqdm(enumerate(input_iter), total=num_files, desc="Encoding images"):
            # Distribute images across GPUs in round-robin fashion
            gpu_id = idx % gpus
            gpu_batches[gpu_id].append((idx, image))
            
            # Process when any GPU batch is full or we've reached the end
            max_batch_size = max(len(gpu_batch) for gpu_batch in gpu_batches)
            if max_batch_size >= batch_size or idx == num_files - 1:
                # Prepare arguments for each GPU
                gpu_args = []
                for gpu_id, gpu_batch in enumerate(gpu_batches):
                    if gpu_batch:  # Only process non-empty batches
                        gpu_args.append((gpu_id, gpu_batch, model_url))
                
                # Process current batches in parallel across GPUs
                if gpu_args:
                    batch_results = pool.map(encode_image_worker, gpu_args)
                    
                    # Flatten results and sort by index