utils/inference.py

import os
import io
import shutil
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

BASE_DIR = Path(__file__).resolve().parents[1]
MODEL_PATH = Path(os.getenv("MODEL_PATH", BASE_DIR / "assets" / "generator.pt"))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def resolve_model_path() -> Path:
    if MODEL_PATH.exists():
        return MODEL_PATH

    repo_id = os.getenv("MODEL_REPO_ID", "amit-saw/gan-sketch-to-image").strip()
    filename = os.getenv("MODEL_FILENAME", MODEL_PATH.name).strip() or MODEL_PATH.name

    if not repo_id:
        raise FileNotFoundError(
            f"Model file not found at {MODEL_PATH}. "
            "Set MODEL_REPO_ID and MODEL_FILENAME to download it from Hugging Face."
        )

    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    downloaded_file = Path(
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            repo_type="model",
            token=os.getenv("HF_TOKEN"),
            local_dir=str(MODEL_PATH.parent),
        )
    )

    if downloaded_file.resolve() != MODEL_PATH.resolve():
        shutil.copyfile(downloaded_file, MODEL_PATH)

    return MODEL_PATH


# Load once at import time so each request is fast.
MODEL = torch.jit.load(str(resolve_model_path()), map_location=DEVICE)
MODEL.eval()

PREPROCESS = transforms.Compose(
    [
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


def tensor_to_rgb_image(tensor: torch.Tensor) -> Image.Image:
    # Model output is in [-1, 1], convert to [0, 255].
    array = tensor.detach().cpu().clamp(-1, 1)
    array = ((array + 1) / 2.0).permute(1, 2, 0).numpy()
    array = (array * 255.0).astype(np.uint8)
    return Image.fromarray(array)


def generate_image_from_sketch_bytes(input_bytes: bytes) -> bytes:
    image = Image.open(io.BytesIO(input_bytes)).convert("RGB")
    input_tensor = PREPROCESS(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        output_tensor = MODEL(input_tensor)[0]

    output_image = tensor_to_rgb_image(output_tensor)
    output_buffer = io.BytesIO()
    output_image.save(output_buffer, format="PNG")
    return output_buffer.getvalue()