File size: 2,362 Bytes
74f0b48 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """
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
|