utils.py

import gradio as gr
from io import BytesIO
from PIL.Image import Image, open as open_image
from os import getenv
import requests
from tempfile import NamedTemporaryFile
import torch


# Try to import spaces decorator (for Hugging Face Spaces), otherwise use no-op decorator.
try:
    import spaces
    spaces_gpu = spaces.GPU
except ImportError:
    # For local development, use a no-op decorator because spaces is not available.
    def spaces_gpu(func):
        return func

def get_pytorch_device() -> str:
    return ("cuda" if torch.cuda.is_available() # Nvidia CUDA and AMD ROCm
       else "xpu" if torch.xpu.is_available() # Intel XPU
       else "mps" if torch.mps.is_available() # Apple Silicon
       else "cpu") # gl bro 🫠

def request_image(url: str) -> Image:
    try:
        response = requests.get(url, timeout=int(getenv("REQUEST_TIMEOUT")))
        response.raise_for_status()
        return open_image(BytesIO(response.content))
    except requests.HTTPError as e:
        raise gr.Error(f"Failed to fetch image from URL because of HTTP error: {e.response.status_code} {e.response.text}")
    except requests.Timeout as e:
        raise gr.Error(f"Failed to fetch image from URL because the request timed out.")
    except requests.RequestException as e:
        raise gr.Error(f"Failed to fetch image from URL: {str(e)}")

def save_image_to_temp_file(image: Image) -> str:
    image_format = image.format if image.format else 'PNG'
    format_extension = image_format.lower() if image_format else 'png'
    temp_file = NamedTemporaryFile(delete=False, suffix=f".{format_extension}")
    temp_path = temp_file.name
    temp_file.close()
    image.save(temp_path, format=image_format)
    return temp_path