Spaces:
Runtime error
Runtime error
File size: 2,723 Bytes
5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 fd34f8f 5ddd413 e35ce27 fd34f8f 5ddd413 | 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 82 83 84 85 86 87 88 89 90 91 92 | """
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
|