scripts/text_to_image.py

"""
Text-to-image for mesh pipeline: generate a single image from prompt (SD 2.1, local).
Uses same SD_MODEL_PATH / HF_TOKEN as skybox_generator.
"""

import os
import time
from pathlib import Path

import torch

from scripts.skybox_generator import (
    _get_hf_token,
    _raise_if_403,
    _resolve_model_path_and_token,
    FALLBACK_MODEL_ID,
)


def get_device() -> str:
    return "cuda" if torch.cuda.is_available() else "cpu"


def text_to_image(
    prompt: str,
    output_dir: str = "outputs",
    size: int = 512,
    seed: int | None = None,
    model_id: str | None = None,
) -> tuple[str, float]:
    """Generate one image from text. Returns (path_to_image, inference_time_sec)."""
    from diffusers import StableDiffusionPipeline

    device = get_device()
    dtype = torch.float16 if device == "cuda" else torch.float32

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    pretrained, token = _resolve_model_path_and_token()
    load_id = model_id or pretrained
    local_only = os.path.isdir(load_id)
    hub_token = token if token is not True else _get_hf_token()
    pipe = None
    last_error = None
    try:
        pipe = StableDiffusionPipeline.from_pretrained(
            load_id,
            torch_dtype=dtype,
            safety_checker=None,
            token=None if local_only else hub_token,
            local_files_only=local_only,
        )
    except Exception as err:
        last_error = err
        _raise_if_403(err)
        if not local_only:
            try:
                pipe = StableDiffusionPipeline.from_pretrained(
                    FALLBACK_MODEL_ID,
                    torch_dtype=dtype,
                    safety_checker=None,
                    token=hub_token,
                )
            except Exception as err2:
                last_error = err2
                _raise_if_403(err2)
    if pipe is None:
        raise RuntimeError(
            "Could not load Stable Diffusion. On Spaces: add HF_TOKEN in Settings → Variables and secrets "
            "(huggingface.co/settings/tokens). Locally: set HF_TOKEN or download the model first."
        ) from last_error
    pipe = pipe.to(device)

    generator = None
    if seed is not None:
        generator = torch.Generator(device=device).manual_seed(seed)

    t0 = time.perf_counter()
    image = pipe(
        prompt=prompt,
        width=size,
        height=size,
        num_inference_steps=50,
        generator=generator,
    ).images[0]
    t1 = time.perf_counter()

    safe_name = "".join(c if c.isalnum() or c in " -_" else "_" for c in prompt)[:50]
    out_path = os.path.join(output_dir, f"mesh_input_{safe_name.strip()}.png")
    image.save(out_path)
    return out_path, t1 - t0