""" models/text_to_image.py ─────────────────────── DiffusersPipelineLoader — wraps any HuggingFace diffusers StableDiffusionPipeline / StableDiffusionXLPipeline. Inputs (run kwargs) ────────────────────────────────────────────────────────────── prompt : str negative_prompt : str (optional) width : int (default from ModelConfig) height : int (default from ModelConfig) guidance_scale : float num_inference_steps : int seed : int (optional, -1 = random) Outputs (returned dict) ────────────────────────────────────────────────────────────── image : PIL.Image (RGB) seed : int (seed actually used) model : str (model_id used) """ from __future__ import annotations import logging import os import random from typing import Any import torch from PIL import Image from models.base_loader import BaseLoader from utils.device import torch_dtype logger = logging.getLogger(__name__) # ─── helpers ────────────────────────────────────────────────────────────────── def _is_sdxl(model_id: str) -> bool: """Heuristic: does this model_id look like an SDXL checkpoint?""" lower = model_id.lower() return any(k in lower for k in ("xl", "sdxl", "turbo", "lightning")) def _hf_token() -> str | None: """Return HF_TOKEN from env if set, else None.""" return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None def _load_pipeline(model_id: str, device: torch.device, dtype: torch.dtype): """ Import and return the appropriate diffusers pipeline class, auto-detecting SD vs SDXL from the model card config. Strategy -------- 1. Try AutoPipelineForText2Image with token (covers auth-gated models). 2. If that raises an OSError/EnvironmentError related to the Hub (network unreachable, metadata fetch failure), retry with local_files_only=True so a previously-cached copy is used instead of hard-failing. 3. Fall back to plain StableDiffusionPipeline on any other AutoPipeline failure (e.g. pipeline_tag mismatch on older model cards). """ from diffusers import ( StableDiffusionPipeline, AutoPipelineForText2Image, ) token = _hf_token() common_kwargs: dict = { "torch_dtype": dtype, "safety_checker": None, "requires_safety_checker": False, "low_cpu_mem_usage": True, } if token: common_kwargs["token"] = token def _try_auto(extra: dict | None = None): kw = {**common_kwargs, **(extra or {})} return AutoPipelineForText2Image.from_pretrained(model_id, **kw) def _try_sd(extra: dict | None = None): kw = {**common_kwargs, **(extra or {})} return StableDiffusionPipeline.from_pretrained(model_id, **kw) def _handle_hub_error(err: Exception, attempt: str) -> None: """ Translate opaque HF Hub errors into actionable messages, then re-raise. RepositoryNotFoundError (404/401) → The model is gated (license agreement required) and no HF_TOKEN was provided, OR the model_id is wrong. diffusers wraps this as EnvironmentError with the "not a valid model identifier" text. Network / cache miss errors → Hub unreachable; retry with local_files_only handled by caller. """ err_str = str(err).lower() if "not a valid model identifier" in err_str or "repository not found" in err_str: raise EnvironmentError( f"\n\n[T2I] Cannot access '{model_id}' on Hugging Face Hub.\n" f"This model requires you to:\n" f" 1. Accept its license at https://huggingface.co/{model_id}\n" f" 2. Generate a HF access token at https://huggingface.co/settings/tokens\n" f" 3. Add it as HF_TOKEN in your HF Space → Settings → Repository secrets\n" f"\nOriginal error: {err}" ) from err pipe = None try: pipe = _try_auto() except (OSError, EnvironmentError) as hub_err: err_str = str(hub_err).lower() # Gated / not found — give actionable message immediately if "not a valid model identifier" in err_str or "repository not found" in err_str: _handle_hub_error(hub_err, "AutoPipeline") # always raises # Network unreachable — retry from local cache if any(k in err_str for k in ("not cached", "fetch metadata", "connection", "network", "offline", "cannot reach", "name or service not known")): logger.warning( "Hub unreachable for %s. Retrying with local_files_only=True …", model_id, ) try: pipe = _try_auto({"local_files_only": True}) logger.info("Loaded %s from local cache.", model_id) except Exception as cache_err: raise OSError( f"Cannot load model {model_id}: Hub unreachable and no local cache found.\n" f" Hub error : {hub_err}\n" f" Cache error: {cache_err}\n" "Tip: run the pipeline once with internet access to warm the cache." ) from cache_err else: raise except Exception as e: logger.warning("AutoPipeline failed (%s); falling back to StableDiffusionPipeline: %s", model_id, e) try: pipe = _try_sd() except (OSError, EnvironmentError) as hub_err2: err_str2 = str(hub_err2).lower() if "not a valid model identifier" in err_str2 or "repository not found" in err_str2: _handle_hub_error(hub_err2, "StableDiffusionPipeline") # always raises if any(k in err_str2 for k in ("not cached", "fetch metadata", "connection", "network", "offline", "cannot reach", "name or service not known")): logger.warning( "Hub unreachable for SD fallback %s. Retrying with local_files_only=True …", model_id, ) try: pipe = _try_sd({"local_files_only": True}) except Exception as cache_err2: raise OSError( f"Cannot load model {model_id}: Hub unreachable and no local cache.\n" f" Hub error : {hub_err2}\n" f" Cache error: {cache_err2}\n" ) from cache_err2 else: raise pipe = pipe.to(device) # Enable memory-efficient attention when available try: pipe.enable_xformers_memory_efficient_attention() logger.info("xformers attention enabled") except Exception: pass try: pipe.enable_attention_slicing() except Exception: pass return pipe # ─── Loader ────────────────────────────────────────────────────────────────── class DiffusersPipelineLoader(BaseLoader): def load(self) -> None: if self._loaded: logger.info("Already loaded — skipping") return logger.info("Loading text-to-image model: %s", self.model_id) dtype = torch_dtype(self.device) self.pipe = _load_pipeline(self.model_id, self.device, dtype) self._loaded = True logger.info("Model ready: %s", self.model_id) def run(self, **inputs: Any) -> dict[str, Any]: if not self._loaded: self.load() prompt: str = inputs.get("prompt", "a photograph of a scene") negative_prompt: str = inputs.get("negative_prompt", "blurry, low quality, distorted") width: int = int(inputs.get("width", self.kwargs.get("width", 512))) height: int = int(inputs.get("height", self.kwargs.get("height", 512))) guidance_scale: float = float(inputs.get("guidance_scale", self.kwargs.get("guidance_scale", 7.5))) num_steps: int = int(inputs.get("num_inference_steps", self.kwargs.get("num_inference_steps", 25))) seed: int = int(inputs.get("seed", -1)) if seed == -1: seed = random.randint(0, 2**32 - 1) generator = torch.Generator(device=self.device).manual_seed(seed) logger.info( "Generating image prompt=%r size=%dx%d steps=%d cfg=%.1f seed=%d", prompt[:80], width, height, num_steps, guidance_scale, seed, ) call_kwargs: dict[str, Any] = { "prompt": prompt, "width": width, "height": height, "num_inference_steps": num_steps, "generator": generator, } # guidance_scale=0 is only valid for SDXL-Turbo if guidance_scale > 0: call_kwargs["guidance_scale"] = guidance_scale if negative_prompt: call_kwargs["negative_prompt"] = negative_prompt result = self.pipe(**call_kwargs) image: Image.Image = result.images[0] return {"image": image, "seed": seed, "model": self.model_id}