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 _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)
    pipe = None
    try:
        pipe = StableDiffusionPipeline.from_pretrained(
            load_id,
            torch_dtype=dtype,
            safety_checker=None,
            token=None if local_only else (token or True),
            local_files_only=local_only,
        )
    except Exception:
        if not local_only:
            try:
                pipe = StableDiffusionPipeline.from_pretrained(
                    FALLBACK_MODEL_ID,
                    torch_dtype=dtype,
                    safety_checker=None,
                    token=token or True,
                )
            except Exception:
                pass
    if pipe is None:
        raise RuntimeError(
            "Could not load Stable Diffusion. Need internet (first run). Set HF_TOKEN if behind firewall."
        )
    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