util.py

import torch

import numpy as np
import random
import os
from typing import Tuple, Optional


def get_prompt_template(mode: str = 'default') -> Tuple[str, int, int]:
    '''
    Generate a prompt template based on the specified mode.

    Args:
        mode (str, optional): The mode for selecting the prompt template. Default is 'default'.

    Returns:
        Tuple[str, int, int]: A tuple containing the generated prompt template, the position of the placeholder '{}',
                             and the length of the prompt.

    Notes:
        If the mode is 'random', a random prompt template is chosen from a predefined list.
    '''
    prompt_template = 'A photo of {}'

    if mode == 'random':
        prompt_templates = [
            'a photo of a {}', 'a photograph of a {}', 'an image of a {}', '{}',
            'a cropped photo of a {}', 'a good photo of a {}', 'a photo of one {}',
            'a bad photo of a {}', 'a photo of the {}', 'a photo of {}', 'a blurry photo of a {}',
            'a picture of a {}', 'a photo of a scene where {}'
        ]
        prompt_template = random.choice(prompt_templates)

    # Calculate prompt length and text position
    prompt_length = 1 + len(prompt_template.split(' ')) + 1 - 1  # eos, sos => 1 + 1, {} => -1
    text_pos_at_prompt = 1 + prompt_template.split(' ').index('{}')

    return prompt_template, text_pos_at_prompt, prompt_length


# Reproducibility
def fix_seed(seed: int = 0) -> None:
    '''
    Set seeds for random number generators to ensure reproducibility.

    Args:
        seed (int, optional): The seed value. Default is 0.
    '''
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # multi-GPU
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)


def seed_worker(worker_id: int) -> None:
    '''
    Set a seed for a worker process to ensure reproducibility in PyTorch DataLoader.

    Args:
        worker_id (int): The ID of the worker process.
    '''
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)