""" pipeline_manager.py ------------------- Loads diffusion pipelines from an editable registry (models.json) and runs generation across multiple base families (SD1.5 / SDXL / FLUX) and multiple input modes (txt2img / img2img / IP-Adapter / Face identity). Designed for Hugging Face ZeroGPU: pipelines are built/cached on CPU and moved to CUDA inside the @spaces.GPU-decorated caller (see app.py). Nothing here calls .cuda() at import time. """ import os import json import gc import hashlib import urllib.request from pathlib import Path import torch from PIL import Image # --------------------------------------------------------------------------- # Constants / paths # --------------------------------------------------------------------------- HERE = Path(__file__).parent REGISTRY_PATH = HERE / "models.json" DOWNLOAD_DIR = Path(os.environ.get("CS_CACHE_DIR", "/tmp/cs_models")) DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", "").strip() HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32 # SD1.5 / SDXL are most stable in float16; FLUX prefers bfloat16. DTYPE_SD = torch.float16 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Modes supported per base family. Used by the UI to gate options. SUPPORTED_MODES = { "sd15": ["txt2img", "img2img", "ip_adapter", "face_id", "pose"], "sdxl": ["txt2img", "img2img", "ip_adapter", "face_id", "pose"], "flux": ["txt2img", "img2img"], } MODE_LABELS = { "txt2img": "Text → Image", "img2img": "Image → Image (denoise)", "ip_adapter": "IP-Adapter (style / subject)", "face_id": "Face identity (FaceID)", "pose": "Pose lock (ControlNet OpenPose)", } # Negative textual-inversion embeddings (the realism trick used on Civitai). # Loaded best-effort in _build_base_pipeline; add the `token` to a model's # negative_prompt to activate it. Repos are public, no token needed. NEG_EMBEDDINGS = { "easynegative": dict(repo="gsdf/EasyNegative", weight="EasyNegative.safetensors", token="EasyNegative"), } def _make_scheduler(pipe, name): """Build a sampler by name (A1111 conventions). Defaults to Euler Ancestral.""" cfg = pipe.scheduler.config if name in ("dpmpp_2m_karras", "dpmpp_2m"): from diffusers import DPMSolverMultistepScheduler return DPMSolverMultistepScheduler.from_config( cfg, algorithm_type="dpmsolver++", use_karras_sigmas=name.endswith("karras")) if name in ("dpmpp_sde_karras", "dpmpp_sde"): from diffusers import DPMSolverSinglestepScheduler return DPMSolverSinglestepScheduler.from_config( cfg, algorithm_type="sde-dpmsolver++", use_karras_sigmas=name.endswith("karras")) if name == "euler": from diffusers import EulerDiscreteScheduler return EulerDiscreteScheduler.from_config(cfg) from diffusers import EulerAncestralDiscreteScheduler return EulerAncestralDiscreteScheduler.from_config(cfg) # --------------------------------------------------------------------------- # Registry # --------------------------------------------------------------------------- def load_registry(): """Read models.json and return the list of enabled model configs.""" with open(REGISTRY_PATH, "r", encoding="utf-8") as f: data = json.load(f) models = [m for m in data.get("models", []) if m.get("enabled", True)] return models def get_model(models, model_id): for m in models: if m["id"] == model_id: return m return None # --------------------------------------------------------------------------- # Thai → English prompt translation (the SD/SDXL/FLUX text encoders are English; # Thai prompts otherwise produce unrelated images). Runs on the Space, no API key. # --------------------------------------------------------------------------- TRANSLATORS = { "nllb": "facebook/nllb-200-distilled-600M", "typhoon": "scb10x/llama3.2-typhoon2-3b-instruct", } _TRANSLATOR_CACHE = {} def has_thai(text): return any("฀" <= ch <= "๿" for ch in (text or "")) # Full-body framing cues — if present, SD1.5's 512x768 canvas crops to a portrait, # so we give the canvas more vertical room (see run_generation). Checked on the # already-translated English prompt; Thai เต็มตัว/ทั้งตัว included as a safety net. _FULL_BODY_CUES = ("full body", "full-body", "head to toe", "head-to-toe", "full length", "full-length", "full shot", "entire body", "whole body", "standing", "เต็มตัว", "ทั้งตัว", "เห็นเท้า") def wants_full_body(text): t = (text or "").lower() return any(c in t for c in _FULL_BODY_CUES) def _load_translator(engine): if engine in _TRANSLATOR_CACHE: return _TRANSLATOR_CACHE[engine] name = TRANSLATORS[engine] if engine == "nllb": from transformers import AutoTokenizer, AutoModelForSeq2SeqLM tok = AutoTokenizer.from_pretrained(name) model = AutoModelForSeq2SeqLM.from_pretrained(name, torch_dtype=DTYPE_SD) else: # typhoon (causal LM) from transformers import AutoTokenizer, AutoModelForCausalLM tok = AutoTokenizer.from_pretrained(name) model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=DTYPE_SD) model.eval() _TRANSLATOR_CACHE[engine] = (tok, model) return tok, model def translate_prompt(text, engine): """Translate a Thai prompt to English. Pass-through if empty/English/off. MUST be called inside the @spaces.GPU context (uses CUDA when available).""" if not text or engine in (None, "off") or not has_thai(text): return text try: tok, model = _load_translator(engine) model = model.to(DEVICE) if engine == "nllb": tok.src_lang = "tha_Thai" inputs = tok(text, return_tensors="pt", truncation=True, max_length=400).to(DEVICE) bos = tok.convert_tokens_to_ids("eng_Latn") out = model.generate(**inputs, forced_bos_token_id=bos, max_new_tokens=256, num_beams=4) return tok.batch_decode(out, skip_special_tokens=True)[0].strip() # typhoon: ask the LLM to rewrite as a clean English image prompt msgs = [ {"role": "system", "content": "You convert Thai text-to-image prompts into " "an English prompt for a PHOTOREALISTIC Stable Diffusion model. Output a " "COMPACT comma-separated list of English tags / short phrases (booru-tag " "style) — NOT full sentences. Omit articles and filler words (a, an, the, " "with, that is). Keep it short to fit a 77-token limit, but INCLUDE EVERY " "detail from the input — especially the location/scene, camera framing " "(e.g. full body), clothing and pose; never drop the setting. Treat it as a " "real candid photograph (natural skin texture, real hair, lifelike eyes, " "natural light). NEVER use illustration/painting/anime/CG words such as " "'masterpiece', 'best quality', 'render', '3d', 'anime' or 'painting'. " "Output ONLY the comma-separated tags — no quotes, no explanation."}, {"role": "user", "content": text}, ] chat = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False) inputs = tok(chat, return_tensors="pt").to(DEVICE) eos = tok.eos_token_id pad = eos[0] if isinstance(eos, (list, tuple)) else eos with torch.no_grad(): out = model.generate(**inputs, max_new_tokens=256, do_sample=False, pad_token_id=pad) gen = out[0][inputs["input_ids"].shape[1]:] return tok.decode(gen, skip_special_tokens=True).strip().strip('"') except Exception as e: # noqa import traceback as _tb print(f"[translate] {engine} failed, using original text: " f"{type(e).__name__}: {e}") _tb.print_exc() return text # --------------------------------------------------------------------------- # Download helpers (Civitai / arbitrary URL → local cache) # --------------------------------------------------------------------------- def _download_url(url): """Download a (Civitai or other) URL to the local cache and return the path.""" if not url: return None fname = hashlib.sha1(url.encode()).hexdigest()[:16] + ".safetensors" dest = DOWNLOAD_DIR / fname if dest.exists() and dest.stat().st_size > 1_000_000: return str(dest) dl_url = url if "civitai.com" in url and CIVITAI_TOKEN and "token=" not in url: sep = "&" if "?" in url else "?" dl_url = f"{url}{sep}token={CIVITAI_TOKEN}" req = urllib.request.Request(dl_url, headers={"User-Agent": "Mozilla/5.0"}) print(f"[download] {url} -> {dest}") with urllib.request.urlopen(req) as resp, open(dest, "wb") as out: while True: chunk = resp.read(1 << 20) if not chunk: break out.write(chunk) # A real model is many MB; a tiny file means Civitai returned a login/redirect page. if dest.stat().st_size < 1_000_000: dest.unlink(missing_ok=True) raise ValueError( "ดาวน์โหลดโมเดลจาก Civitai ไม่สำเร็จ — โมเดลนี้ต้องตั้งค่า CIVITAI_TOKEN " "ใน Space Settings → Variables and secrets ก่อน / Civitai download failed: " "set CIVITAI_TOKEN in the Space secrets to use this model." ) return str(dest) # --------------------------------------------------------------------------- # Pipeline cache # --------------------------------------------------------------------------- # Keyed by model id. Stores the base txt2img pipeline (CPU). Adapters are loaded # on demand and tracked via the `_cs_adapter` attribute on the pipe. _PIPE_CACHE = {} _FACE_APP = None # lazy insightface FaceAnalysis def _free_cache(keep_id=None): """Evict cached pipelines except keep_id to bound memory (simple LRU-ish).""" for k in list(_PIPE_CACHE.keys()): if k != keep_id: del _PIPE_CACHE[k] gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def _build_base_pipeline(cfg): """Construct the txt2img pipeline for a model config (on CPU).""" base = cfg["base"] common = dict(token=HF_TOKEN) # Some checkpoint merges overflow to NaN in fp16 (rainbow-noise output); # such models set "dtype": "fp32" in the registry. dt = torch.float32 if cfg.get("dtype") == "fp32" else DTYPE_SD if base == "sd15": from diffusers import StableDiffusionPipeline if cfg.get("single_file_url"): local = _download_url(cfg["single_file_url"]) pipe = StableDiffusionPipeline.from_single_file( local, torch_dtype=dt, safety_checker=None ) else: pipe = StableDiffusionPipeline.from_pretrained( cfg["repo_id"], torch_dtype=dt, safety_checker=None, **common ) elif base == "sdxl": from diffusers import StableDiffusionXLPipeline if cfg.get("single_file_url"): local = _download_url(cfg["single_file_url"]) pipe = StableDiffusionXLPipeline.from_single_file(local, torch_dtype=dt) else: pipe = StableDiffusionXLPipeline.from_pretrained( cfg["repo_id"], torch_dtype=dt, **common ) elif base == "flux": from diffusers import FluxPipeline pipe = FluxPipeline.from_pretrained(cfg["repo_id"], torch_dtype=DTYPE, **common) else: raise ValueError(f"Unknown base family: {base}") # Apply LoRA if this entry is a LoRA model. if cfg.get("type") == "lora": scale = float(cfg.get("lora_scale", 0.8)) # Resolve to a local .safetensors path (HF repo or direct/Civitai URL). if cfg.get("lora_repo_id"): from huggingface_hub import hf_hub_download local = hf_hub_download(cfg["lora_repo_id"], cfg["lora_weight_name"]) \ if cfg.get("lora_weight_name") else None if local is None: pipe.load_lora_weights(cfg["lora_repo_id"]) local = "__loaded__" else: local = _download_url(cfg.get("lora_url")) if local and local != "__loaded__": try: pipe.load_lora_weights(local) except Exception as e: # noqa # Some Civitai/kohya LoRAs carry text-encoder keys diffusers can't # convert ("list index out of range"). Retry with UNet-only keys — # the UNet holds most of the character/style effect. print(f"[lora] full load failed ({e}); retrying UNet-only") from safetensors.torch import load_file sd = load_file(local) sd = {k: v for k, v in sd.items() if not k.startswith("lora_te")} pipe.load_lora_weights(sd) try: pipe.fuse_lora(lora_scale=scale) except Exception as e: # noqa print(f"[lora] fuse skipped: {e}") # Optional VAE override (known-good VAE for models with a broken one). if cfg.get("vae"): from diffusers import AutoencoderKL pipe.vae = AutoencoderKL.from_pretrained(cfg["vae"], torch_dtype=dt) # SD1.5 / SDXL community checkpoints match the A1111 / ComfyUI look best with # a community sampler. Per-model `sampler` in the registry (e.g. dpmpp_2m_karras # for photoreal); defaults to Euler Ancestral. if base in ("sd15", "sdxl"): pipe.scheduler = _make_scheduler(pipe, cfg.get("sampler", "euler_a")) # Negative textual-inversion embeddings (realism boost). Best-effort: a failed # download just means the token falls through as plain words in the negative. for emb in cfg.get("neg_embeddings", []): spec = NEG_EMBEDDINGS.get(emb) if not spec: continue try: pipe.load_textual_inversion(spec["repo"], weight_name=spec["weight"], token=spec["token"]) except Exception as e: # noqa print(f"[ti] negative embedding '{emb}' failed: {type(e).__name__}: {e}") pipe.set_progress_bar_config(disable=True) pipe._cs_adapter = None # track loaded IP-Adapter / FaceID state return pipe def get_pipeline(cfg): """Return a cached base pipeline for the model, building it if needed.""" mid = cfg["id"] if mid not in _PIPE_CACHE: _free_cache(keep_id=None) # one big model at a time on ZeroGPU print(f"[pipeline] building {mid} ({cfg['base']})") _PIPE_CACHE[mid] = _build_base_pipeline(cfg) return _PIPE_CACHE[mid] # --------------------------------------------------------------------------- # Adapter management (IP-Adapter / FaceID) # --------------------------------------------------------------------------- _IP_ADAPTER_SPECS = { "sd15": { "ip_adapter": dict(repo="h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin"), "face_id": dict(repo="h94/IP-Adapter-FaceID", subfolder=None, weight_name="ip-adapter-faceid_sd15.bin", image_encoder_folder=None), }, "sdxl": { "ip_adapter": dict(repo="h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter-plus_sdxl_vit-h.bin"), "face_id": dict(repo="h94/IP-Adapter-FaceID", subfolder=None, weight_name="ip-adapter-faceid_sdxl.bin", image_encoder_folder=None), }, } def _ensure_adapter(pipe, base, mode): """Load the right IP-Adapter for `mode`, unloading any previous one.""" want = mode if mode in ("ip_adapter", "face_id") else None if pipe._cs_adapter == want: return try: pipe.unload_ip_adapter() except Exception: pass pipe._cs_adapter = None if want is None: return spec = _IP_ADAPTER_SPECS[base][want] kwargs = {k: v for k, v in spec.items() if k != "repo"} pipe.load_ip_adapter(spec["repo"], **kwargs) pipe._cs_adapter = want def _get_face_app(): global _FACE_APP if _FACE_APP is None: from insightface.app import FaceAnalysis app = FaceAnalysis(name="buffalo_l", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) app.prepare(ctx_id=0, det_size=(640, 640)) _FACE_APP = app return _FACE_APP def _face_embeds(image): """Return a torch tensor of FaceID embeddings for the largest face.""" import numpy as np import cv2 app = _get_face_app() arr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR) faces = app.get(arr) if not faces: raise ValueError("ไม่พบใบหน้าในรูปต้นแบบ / No face detected in the reference image.") faces = sorted(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1])) emb = torch.from_numpy(faces[-1].normed_embedding) # [512] # diffusers IP-Adapter-FaceID expects [2, 1, 1, 512]: [neg, pos] for CFG. emb = emb.unsqueeze(0).unsqueeze(0).unsqueeze(0) # [1, 1, 1, 512] return torch.cat([torch.zeros_like(emb), emb], dim=0).to(DTYPE_SD) # --------------------------------------------------------------------------- # Generation # --------------------------------------------------------------------------- # --------------------------------------------------------------------------- # ControlNet (OpenPose) — locks the generated subject to an uploaded pose. # --------------------------------------------------------------------------- _CONTROLNET = {} _OPENPOSE = None def _get_controlnet(base): if base in _CONTROLNET: return _CONTROLNET[base] from diffusers import ControlNetModel repos = { "sd15": "lllyasviel/control_v11p_sd15_openpose", "sdxl": "xinsir/controlnet-openpose-sdxl-1.0", } if base not in repos: raise ValueError("Pose (ControlNet) รองรับ SD1.5 / SDXL เท่านั้น.") cn = ControlNetModel.from_pretrained(repos[base], torch_dtype=DTYPE_SD) _CONTROLNET[base] = cn return cn def _get_openpose(): global _OPENPOSE if _OPENPOSE is None: from controlnet_aux import OpenposeDetector _OPENPOSE = OpenposeDetector.from_pretrained("lllyasviel/Annotators") return _OPENPOSE def _safe_call(pipe_obj, call): """Run the pipeline; if clip_skip trips a version incompatibility, retry without it.""" try: return pipe_obj(**call).images[0] except (AttributeError, TypeError) as e: if "clip_skip" in call: print(f"[clip_skip] disabled for this run due to: {e}") call.pop("clip_skip", None) return pipe_obj(**call).images[0] raise def run_generation(cfg, mode, prompt, negative_prompt, ref_image, steps, guidance, denoise, ip_scale, width, height, seed): """Run one generation. MUST be called inside a @spaces.GPU context.""" base = cfg["base"] if mode not in SUPPORTED_MODES[base]: raise ValueError( f"โหมด '{MODE_LABELS.get(mode, mode)}' ใช้กับ base {base.upper()} ไม่ได้ " f"(รองรับ: {', '.join(MODE_LABELS[m] for m in SUPPORTED_MODES[base])})" ) pipe = get_pipeline(cfg) pipe = pipe.to(DEVICE) generator = None if seed is not None and int(seed) >= 0: generator = torch.Generator(device=DEVICE).manual_seed(int(seed)) # Front-load a per-model photographic style prefix (e.g. "RAW photo, film grain") # so the strongest realism cue survives CLIP's 77-token truncation. Applied AFTER # translation (run_generation gets the already-English prompt), so it reaches the # model verbatim regardless of the translator. Then trigger, then the user prompt. _parts = [cfg.get("style_prefix"), cfg.get("trigger"), prompt] full_prompt = ", ".join(p.strip() for p in _parts if p and str(p).strip()).strip(", ") # Full-body framing fix: on SD1.5 a 512x768 canvas crops standing/seated subjects # to a portrait even when "full body" is requested. Give it more vertical room. if base == "sd15" and mode == "txt2img" and wants_full_body(prompt): height = max(int(height), 896) call = dict( prompt=full_prompt, num_inference_steps=int(steps), generator=generator, width=int(width), height=int(height), ) # FLUX uses `guidance_scale` differently and has no negative prompt. if base == "flux": call["guidance_scale"] = float(guidance) else: call["guidance_scale"] = float(guidance) call["negative_prompt"] = negative_prompt or None # ----- mode wiring ----- if mode == "txt2img": _ensure_adapter(pipe, base, None) # Hires fix: generate at base resolution, then upscale + low-denoise # refine pass. This is the single biggest realism lever for SD1.5 — # it adds skin pores / hair strands / sharpness the native pass lacks. hires = cfg.get("hires") if hires and base in ("sd15", "sdxl"): base_img = _safe_call(pipe, call) scale = float(hires.get("scale", 1.5)) hw = max(8, int(width * scale) // 8 * 8) hh = max(8, int(height * scale) // 8 * 8) up = base_img.resize((hw, hh), Image.LANCZOS) from diffusers import AutoPipelineForImage2Image i2i = AutoPipelineForImage2Image.from_pipe(pipe).to(DEVICE) gen2 = None if seed is not None and int(seed) >= 0: gen2 = torch.Generator(device=DEVICE).manual_seed(int(seed)) hcall = dict( prompt=full_prompt, negative_prompt=negative_prompt or None, image=up, strength=float(hires.get("denoise", 0.4)), num_inference_steps=int(hires.get("steps", steps)), guidance_scale=float(guidance), generator=gen2, ) return _safe_call(i2i, hcall) elif mode == "img2img": _ensure_adapter(pipe, base, None) if base != "flux" else None if ref_image is None: raise ValueError("img2img ต้องอัปโหลดรูปต้นแบบก่อน / Upload a reference image first.") from diffusers import AutoPipelineForImage2Image i2i = AutoPipelineForImage2Image.from_pipe(pipe).to(DEVICE) call.pop("width"); call.pop("height") call["image"] = ref_image.convert("RGB") call["strength"] = float(denoise) return _safe_call(i2i, call) elif mode == "ip_adapter": if ref_image is None: raise ValueError("IP-Adapter ต้องอัปโหลดรูปต้นแบบก่อน / Upload a reference image first.") _ensure_adapter(pipe, base, "ip_adapter") pipe.set_ip_adapter_scale(float(ip_scale)) call["ip_adapter_image"] = ref_image.convert("RGB") elif mode == "face_id": if ref_image is None: raise ValueError("Face identity ต้องอัปโหลดรูปใบหน้าก่อน / Upload a face image first.") _ensure_adapter(pipe, base, "face_id") pipe.set_ip_adapter_scale(float(ip_scale)) embeds = _face_embeds(ref_image).to(DEVICE) call["ip_adapter_image_embeds"] = [embeds] elif mode == "pose": if ref_image is None: raise ValueError("Pose ต้องอัปโหลดรูปท่าทางก่อน / Upload a pose reference image first.") _ensure_adapter(pipe, base, None) detector = _get_openpose() pose_img = detector(ref_image.convert("RGB")).resize((int(width), int(height))) cn = _get_controlnet(base).to(DEVICE) if base == "sdxl": from diffusers import StableDiffusionXLControlNetPipeline cn_pipe = StableDiffusionXLControlNetPipeline.from_pipe(pipe, controlnet=cn).to(DEVICE) else: from diffusers import StableDiffusionControlNetPipeline cn_pipe = StableDiffusionControlNetPipeline.from_pipe(pipe, controlnet=cn).to(DEVICE) call["image"] = pose_img return _safe_call(cn_pipe, call) return _safe_call(pipe, call)