utils/misc.py

# misc.py file contains miscellaneous utility functions
import math
import sys
import logging
import os
import subprocess
from utils.constants import TMPDIR, MAX_SEED
import numpy as np

def get_output_name(input_image=None, output_image=None, overlay_image=None, bordered_image_output=None):
    """
    Check input_image, output_image, overlay_image, bordered_image_output to find a filename
    other than None or generic 'image.PNG' and return the filename without extension to use as output_name
    
    Args:
        input_image (str): Path to input image
        output_image (str): Path to output image
        overlay_image (str): Path to overlay image
        bordered_image_output (str): Path to bordered image output
        
    Returns:
        str: Base filename without extension to use as output name, or 'hexagrid_output' if none found
    """
    # Check each input in priority order
    for image_path in [bordered_image_output, overlay_image, output_image, input_image]:
        if image_path and isinstance(image_path, str) and get_filename_from_path(image_path.lower()) != "image.png":
            # Get the filename without path or extension
            base_name = get_filename_from_path(image_path)
            name_without_ext = os.path.splitext(base_name)[0]
            
            # Filter out temporary filenames that start with tmp or have random hex strings
            if not (name_without_ext.startswith("tmp") or 
                   (len(name_without_ext) >= 8 and all(c in '0123456789abcdef' for c in name_without_ext))):
                return name_without_ext
    
    # Default name if no suitable filename found
    return "hexagrid_output"

def get_filename_from_path(file_path):
    """Extract filename from file path."""
    if file_path is None:
        return ""
    import os
    return os.path.basename(file_path)

def pause():
    """
    Pauses the execution until any key is pressed.
    """
    if sys.platform.startswith('win'):
        import msvcrt
        print("Press any key to continue...")
        msvcrt.getch()
    else:
        import termios
        import tty
        print("Press any key to continue...")
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        try:
            tty.setraw(fd)
            sys.stdin.read(1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)


def install(package):
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

def get_filename(file):
    filename = None
    if file is not None:
        filename = file.name
    return filename

def get_extension(file):
    extension = None
    if file is not None:
        extension = file.name.split(".")[-1]
    return extension


def convert_ratio_to_dimensions(ratio, height=512, rotate90=False):
    """
    Calculate pixel dimensions based on a given aspect ratio and base height.

    This function computes the width and height in pixels for an image, ensuring that both dimensions are divisible by 16. The height is adjusted upwards to the nearest multiple of 16 if necessary, and the width is calculated based on the adjusted height and the provided aspect ratio. Additionally, it ensures that both width and height are at least 16 pixels to avoid extremely small dimensions.

    Parameters:
        ratio (float): The aspect ratio of the image (width divided by height).
        height (int, optional): The base height in pixels. Defaults to 512.

    Returns:
        tuple: A tuple containing the calculated (width, height) in pixels, both divisible by 16.
    """
    base_height = 512
    # Scale the height based on the provided height parameter
    # Ensure the height is at least base_height
    scaled_height = max(height, base_height)
    # Adjust the height to be divisible by 16
    adjusted_height = math.ceil(scaled_height / 16) * 16
    # Calculate the width based on the ratio
    calculated_width = int(adjusted_height * ratio)
    # Adjust the width to be divisible by 16
    adjusted_width = math.ceil(calculated_width / 16) * 16
    if rotate90:
        adjusted_width, adjusted_height = adjusted_height, adjusted_width
    return adjusted_width, adjusted_height

def update_dimensions_on_ratio(aspect_ratio_str, base_height):
    # Convert aspect_ratio from a string in format "W:H" into numbers and compute new dimensions.
    width_ratio, height_ratio = map(int, aspect_ratio_str.split(":"))
    aspect_ratio = width_ratio / height_ratio
    new_width, new_height = convert_ratio_to_dimensions(aspect_ratio, base_height)
    return new_width, new_height

def install_torch():    
    print("\nInstalling PyTorch with CUDA support...")
    # Define the package and index URL
    package = "torch==2.4.0"
    index_url = "https://download.pytorch.org/whl/cu124"
    # Construct the pip install command
    command = [
        "pip", "install", "--force-reinstall", 
        f"{package}", "--index-url", f"{index_url}"
    ]
    # Run the command using subprocess
    subprocess.run(command, check=True)
    print("\nPyTorch installation completed.")
    print("\nInstalling torchvision...")
    package = "torchvision==0.19.0"
    index_url = "https://download.pytorch.org/whl/cu124"
    # Construct the pip install command
    command = [
        "pip", "install", "--force-reinstall", 
        f"{package}", "--index-url", f"{index_url}"
    ]
    # Run the command using subprocess
    subprocess.run(command, check=True)
    print("\nPlease restart the kernel to use the newly installed PyTorch version.")

def _get_output(cmd):
    try:
        return subprocess.check_output(cmd).decode("utf-8")
    except Exception as ex:
        logging.exception(ex)
    return None

def install_cuda_toolkit():    
    #CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
    # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run"
    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
    print("\nDownloading CUDA Toolkit from %s" % CUDA_TOOLKIT_URL)
    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])    
    os.environ["CUDA_HOME"] = "/usr/local/cuda"
    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
        os.environ["CUDA_HOME"],
        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
    )
    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
    print("\nPlease restart the kernel to use the newly installed CUDA Toolkit.")

def setup_runtime_env():
    from torch import cuda
    logging.info("Python Version: %s" % _get_output(["python", "--version"]))
    logging.info("CUDA Version: %s" % _get_output(["nvcc", "--version"]))
    logging.info("GCC Version: %s" % _get_output(["gcc", "--version"]))
    logging.info("CUDA is available: %s" % cuda.is_available())
    logging.info("CUDA Device Capability: %s" % (cuda.get_device_capability(),))

    # Install Pre-compiled CUDA extensions (Fallback to this solution on 12/31/24)
    # Ref: https://huggingface.co/spaces/zero-gpu-explorers/README/discussions/110
    ##ext_dir = os.path.join(os.path.dirname(__file__), "wheels")
    ##for e in os.listdir(ext_dir):
    ##    logging.info("Installing Extensions from %s" % e)
    ##    subprocess.call(
    ##        ["pip", "install", os.path.join(ext_dir, e)], stderr=subprocess.STDOUT
    ##    )
    # Compile CUDA extensions
    # Update on 12/31/24: No module named 'torch'. But it is installed and listed by `pip list`
    # ext_dir = os.path.join(os.path.dirname(__file__), "citydreamer", "extensions")
    # for e in os.listdir(ext_dir):
    #     if os.path.isdir(os.path.join(ext_dir, e)):
    #         subprocess.call(["pip", "install", "."], cwd=os.path.join(ext_dir, e))

    #logging.info("Installed Python Packages: %s" % _get_output(["pip", "list"]))


def get_seed(randomize_seed: bool, seed: int) -> int:
    """
    Get the random seed.
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed

def number_to_letter(n: int, upper_case: bool = True):
    result = ''
    a_char = 97
    if upper_case:
        a_char -= 32
    while abs(n) > 0:
        n, remainder = divmod(abs(n) - 1, 26)
        result = chr(a_char + remainder) + result
    return result