utils.py

import random
from typing import Tuple

import torch


def prepare_generator(seed: int) -> Tuple[torch.Generator, int]:
    final_seed = seed
    if final_seed is None or final_seed < 0:
        final_seed = random.randint(0, 2**31 - 1)
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(final_seed)
    return generator, final_seed


def sanitize_dimensions(width: int, height: int, max_width: int, max_height: int) -> Tuple[int, int]:
    width = max(256, min(max_width, int(width // 8 * 8)))
    height = max(256, min(max_height, int(height // 8 * 8)))
    return width, height