| """ |
| 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 |
|
|
