Add Gradio app for Spaces

demo.py

@@ -0,0 +1,916 @@
+# =========================
+# ONE-CELL: SDXL + CritiCore + SpecFusion + Gradio UI
+# - Fixes: api_name=False + show_api=False, share=True, VARIANT_* before Blocks
+# - UI style follows your reference (ui_run_once + _run_async + gallery/meta)
+# =========================
+
+import os, re, io, json, time, base64, asyncio, inspect, traceback
+from pathlib import Path
+from typing import List, Dict, Optional, Tuple, Iterable, Set
+
+import torch
+from PIL import Image
+import nest_asyncio
+nest_asyncio.apply()
+
+import gradio as gr
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    DPMSolverMultistepScheduler,
+)
+
+# =========================
+# 0) Variant names / display order  (MUST be BEFORE Blocks)
+# =========================
+VARIANT_LABELS = {
+    "base_multi_llm":                     "1_base_multi_llm",
+    "criticore_on_multi_llm__specfusion": "2_criticore_on_multi_llm__specfusion",
+    "base_original":                      "3_base_original",
+    "criticore_on_original__specfusion":  "4_criticore_on_original__specfusion",
+}
+VARIANT_ORDER = [
+    "1_base_multi_llm",
+    "2_criticore_on_multi_llm__specfusion",
+    "3_base_original",
+    "4_criticore_on_original__specfusion",
+]
+RHO_T_DEFAULT = 0.85  # fixed as requested
+
+# ---- SAFETY: do NOT hardcode API keys ----
+TOGETHER_API_KEY = os.environ.get("TOGETHER_API_KEY", "").strip()
+if not TOGETHER_API_KEY:
+    print("[Warn] TOGETHER_API_KEY is not set. CritiCore will fail if you use Together models.")
+
+# =========================
+# 1) SDXL init
+# =========================
+DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = torch.device(DEVICE_STR)
+DTYPE  = torch.float16 if torch.cuda.is_available() else torch.float32
+SDXL_ID = os.environ.get("SDXL_ID", "stabilityai/stable-diffusion-xl-base-1.0")
+
+print(f"[Init] DEVICE={DEVICE_STR} DTYPE={DTYPE} SDXL_ID={SDXL_ID}")
+
+SDXL_base = StableDiffusionXLPipeline.from_pretrained(SDXL_ID, torch_dtype=DTYPE).to(DEVICE)
+SDXL_i2i  = StableDiffusionXLImg2ImgPipeline.from_pretrained(SDXL_ID, torch_dtype=DTYPE).to(DEVICE)
+
+for p in (SDXL_base, SDXL_i2i):
+    try:
+        p.enable_vae_slicing()
+        p.enable_attention_slicing()
+    except Exception:
+        pass
+    p.scheduler = DPMSolverMultistepScheduler.from_config(p.scheduler.config, use_karras_sigmas=True)
+
+DEFAULT_NEG = (
+    "blurry, low quality, artifacts, watermark, extra fingers, missing limbs, "
+    "over-sharpened, harsh lighting, oversaturated"
+)
+
+@torch.no_grad()
+def decode_image_sdxl(latents: torch.Tensor, pipe: StableDiffusionXLImg2ImgPipeline, output_type="pil"):
+    vae = pipe.vae
+    needs_upcast = (vae.dtype in (torch.float16, torch.bfloat16)) and bool(getattr(vae.config, "force_upcast", False))
+    if needs_upcast:
+        try:
+            pipe.upcast_vae()
+        except Exception:
+            pipe.vae = pipe.vae.to(torch.float32)
+        vae = pipe.vae
+
+    lat = latents.to(device=vae.device, dtype=(next(vae.post_quant_conv.parameters()).dtype))
+    lat = lat / vae.config.scaling_factor
+    out = vae.decode(lat)
+    x = out[0] if isinstance(out, (list, tuple)) else (out.sample if hasattr(out, "sample") else out)
+    if getattr(pipe, "watermark", None) is not None:
+        x = pipe.watermark.apply_watermark(x)
+    img = pipe.image_processor.postprocess(x.detach(), output_type=output_type)[0]
+    return img
+
+@torch.no_grad()
+def base_sample_latent(prompt: str, seed: int = 2025, H: int = 1024, W: int = 1024, neg: str = ""):
+    g = torch.Generator(device=DEVICE).manual_seed(int(seed))
+    out = SDXL_base(
+        prompt=prompt,
+        negative_prompt=neg,
+        height=int(H), width=int(W),
+        guidance_scale=4.5,
+        num_inference_steps=50,
+        generator=g,
+        output_type="latent"
+    )
+    z0 = out.images
+    x0 = decode_image_sdxl(z0, SDXL_i2i)
+    return z0, x0
+
+@torch.no_grad()
+def img2img_latent(prompt: str, image_or_latent, strength: float, guidance: float, steps: int, seed: int):
+    g = torch.Generator(device=DEVICE).manual_seed(int(seed))
+    out = SDXL_i2i(
+        prompt=prompt,
+        image=image_or_latent,
+        strength=float(strength),
+        guidance_scale=float(guidance),
+        num_inference_steps=int(steps),
+        generator=g,
+        output_type="latent",
+        negative_prompt=DEFAULT_NEG
+    )
+    return out.images
+
+def strength_for_last_k(k: int, total_steps: int) -> float:
+    k = max(1, int(k))
+    return min(0.95, max(0.01, float(k) / float(max(1, total_steps))))
+
+# =========================
+# 2) CLIP-77 + text utils
+# =========================
+try:
+    from transformers import CLIPTokenizerFast
+    _clip_tok = CLIPTokenizerFast.from_pretrained("openai/clip-vit-large-patch14")
+    def _count_tokens(txt: str) -> int:
+        return len(_clip_tok(txt, add_special_tokens=True, truncation=False)["input_ids"])
+except Exception:
+    _clip_tok = None
+    def _count_tokens(txt: str) -> int:
+        return int(len(re.findall(r"\w+", txt)) * 1.3)
+
+def _cleanup_commas(s: str) -> str:
+    s = re.sub(r"\s*,\s*", ", ", (s or "").strip())
+    s = re.sub(r"(,\s*){2,}", ", ", s)
+    return s.strip(" ,")
+
+def clip77_strict(text: str, max_tok: int = 77) -> str:
+    text = (text or "").strip()
+    if _count_tokens(text) <= max_tok:
+        return text
+    words = text.split()
+    lo, hi, best = 0, len(words), ""
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        cand = " ".join(words[:mid]) if mid > 0 else ""
+        if _count_tokens(cand) <= max_tok:
+            best = cand; lo = mid + 1
+        else:
+            hi = mid - 1
+    return best.strip()
+
+def _split_tags(s: str) -> List[str]:
+    return [p.strip() for p in re.split(r",|\n", (s or "").strip()) if p.strip()]
+
+def _dedup_keep_order(items: List[str]) -> List[str]:
+    seen, out = set(), []
+    for t in items:
+        key = re.sub(r"\s+", " ", t.lower()).strip()
+        if key and key not in seen:
+            seen.add(key); out.append(t.strip())
+    return out
+
+def _order_tags(subject_first: List[str], rest: List[str]) -> List[str]:
+    buckets = {"subject": [], "style": [], "composition": [], "lighting": [], "color": [], "detail": [], "other": []}
+    style_kw = ("style","painterly","illustration","photorealistic","neon","poster","matte painting","watercolor","cyberpunk")
+    comp_kw  = ("composition","rule of thirds","centered","symmetry","balanced composition")
+    light_kw = ("lighting","light","glow","glowing","rim","sunset","sunrise","golden hour","global illumination","cinematic")
+    color_kw = ("color","palette","vibrant","muted","monochrome","pastel","warm","cool","balanced contrast")
+    detail_kw= ("detailed","hyperdetailed","texture","intricate","high detail","highly detailed","sharp focus","uhd","8k")
+
+    for t in subject_first:
+        if t: buckets["subject"].append(t)
+    for t in rest:
+        lt = t.lower()
+        if   any(k in lt for k in style_kw):  buckets["style"].append(t)
+        elif any(k in lt for k in comp_kw):   buckets["composition"].append(t)
+        elif any(k in lt for k in light_kw):  buckets["lighting"].append(t)
+        elif any(k in lt for k in color_kw):  buckets["color"].append(t)
+        elif any(k in lt for k in detail_kw): buckets["detail"].append(t)
+        else:                                 buckets["other"].append(t)
+
+    return buckets["subject"] + buckets["style"] + buckets["composition"] + buckets["lighting"] + buckets["color"] + buckets["detail"] + buckets["other"]
+
+def pil_to_base64(img: Image.Image, fmt: str = "PNG") -> str:
+    buf = io.BytesIO()
+    img.save(buf, format=fmt)
+    return base64.b64encode(buf.getvalue()).decode("ascii")
+
+async def _maybe_close_async_together(client) -> None:
+    try:
+        if hasattr(client, "aclose") and inspect.iscoroutinefunction(client.aclose):
+            await client.aclose()
+        elif hasattr(client, "close"):
+            fn = client.close
+            if inspect.iscoroutinefunction(fn):
+                await fn()
+            else:
+                try: fn()
+                except Exception: pass
+    except Exception:
+        pass
+
+# =========================
+# 3) Async runner (reference style, but made safer for notebook/gradio)
+# =========================
+def _run_async(coro):
+    """
+    Robust sync->async bridge for:
+      - notebook (loop already running, same thread) via nest_asyncio + run_until_complete
+      - normal python (no running loop) via asyncio.run
+    """
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            # nest_asyncio makes this workable in notebook main loop
+            return loop.run_until_complete(coro)
+        return loop.run_until_complete(coro)
+    except RuntimeError:
+        return asyncio.run(coro)
+
+# =========================
+# 4) CritiCore (Together)
+# =========================
+from together import AsyncTogether
+
+AGGREGATOR_MODEL = os.environ.get("AGGREGATOR_MODEL", "Qwen/Qwen2.5-72B-Instruct-Turbo")
+LLM_MULTI_CANDIDATES = [
+    "meta-llama/Llama-3.3-70B-Instruct-Turbo",
+    "Qwen/Qwen2.5-72B-Instruct-Turbo",
+    "Qwen/Qwen2.5-Coder-32B-Instruct",
+    "deepseek-ai/DeepSeek-V3",
+    "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+]
+_env_list = [s.strip() for s in os.environ.get("VLM_MOA_CANDIDATES","").split(",") if s.strip()]
+VLM_CANDIDATES = _env_list or ["meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"]
+
+TAG_PRESETS = {
+    "hq_preference": {
+        "seed_pos": [
+            "balanced composition",
+            "natural color palette","vibrant colors","balanced contrast",
+            "high detail","highly detailed","hyperdetailed","sharp focus",
+            "UHD","8k"
+        ],
+        "seed_neg": [
+            "low quality","blurry","watermark","jpeg artifacts","overexposed","underexposed",
+            "color banding","extra fingers","missing limbs","disfigured","mutated hands"
+        ]
+    }
+}
+
+_DECOMP_SYS = (
+    "Decompose the user's visual instruction into 3-6 concrete, checkable visual components "
+    "(entities + interactions + spatial relations). Return ONLY JSON: "
+    '{"components":["..."]}'
+)
+
+_TXT_SYS = (
+    "Expand a VERY SHORT visual idea into a COMMA-SEPARATED TAG LIST for SDXL.\n"
+    "Constraints:\n"
+    "- Start with the subject phrase first.\n"
+    "- Prioritize composition, lighting, color, and detail over style.\n"
+    "- Use at most TWO style tags if any.\n"
+    "- 16–26 concise tags total. Commas only, no sentences, no 'and'. No trailing period.\n"
+    "- Prefer human-preference aesthetics; keep 'high detailed', 'sharp focus', '8k', 'UHD'."
+)
+
+def _TAG_RE(tag: str):
+    return re.compile(rf"<\s*{tag}\s*>(.*?)</\s*{tag}\s*>", re.S|re.I)
+
+def _extract_tag(text: str, tag: str, fallback: str = "") -> str:
+    s = (text or "").strip()
+    r = _TAG_RE(tag); m = r.search(s)
+    if m: return m.group(1).strip()
+    s2 = s.replace("&lt;","<").replace("&gt;",">")
+    m2 = r.search(s2)
+    return m2.group(1).strip() if m2 else fallback.strip()
+
+def _summarize_issues_lines(text: str, max_lines: int = 5) -> str:
+    if not text: return ""
+    parts = [p.strip(" -•\t") for p in re.split(r"[\n;]+", text) if p.strip()]
+    parts = parts[:max_lines]
+    return "\n".join(f"- {p}" for p in parts)
+
+class CritiCore:
+    def __init__(self, preset: str = "hq_preference", aggregator_model: str = AGGREGATOR_MODEL):
+        if not os.environ.get("TOGETHER_API_KEY"):
+            raise RuntimeError("Missing TOGETHER_API_KEY in environment.")
+        self.preset = preset
+        self.aggregator = aggregator_model
+
+    async def decompose_components(self, user_prompt: str) -> List[str]:
+        client = AsyncTogether(api_key=os.environ["TOGETHER_API_KEY"])
+        try:
+            tasks = [client.chat.completions.create(
+                model=m,
+                messages=[{"role":"system","content": _DECOMP_SYS},
+                          {"role":"user","content": user_prompt}],
+                temperature=0.4, max_tokens=256
+            ) for m in LLM_MULTI_CANDIDATES]
+            rs = await asyncio.gather(*tasks, return_exceptions=True)
+            texts = []
+            for r in rs:
+                try: texts.append(r.choices[0].message.content)
+                except Exception: pass
+            if not texts: return []
+            joined = "\n\n---\n\n".join(texts)
+            merged = await client.chat.completions.create(
+                model=self.aggregator,
+                messages=[{"role":"system","content": "Merge JSON candidates and return ONLY {'components':[...]}."},
+                          {"role":"user","content": joined}],
+                temperature=0.2, max_tokens=256
+            )
+            txt = merged.choices[0].message.content
+            try:
+                obj = json.loads(txt)
+            except Exception:
+                s,e = txt.find("{"), txt.rfind("}")
+                obj = json.loads(txt[s:e+1]) if (s!=-1 and e!=-1) else {"components":[]}
+            comps = [c.strip() for c in obj.get("components", []) if isinstance(c, str) and c.strip()]
+            return comps[:6]
+        finally:
+            await _maybe_close_async_together(client)
+
+    async def make_tags(self, user_prompt: str, clip77: bool = True) -> Tuple[str, str]:
+        client = AsyncTogether(api_key=os.environ["TOGETHER_API_KEY"])
+        seed = TAG_PRESETS.get(self.preset, TAG_PRESETS["hq_preference"])
+        seed_pos = _dedup_keep_order(seed["seed_pos"])
+        seed_neg = seed["seed_neg"]
+        try:
+            tasks = [client.chat.completions.create(
+                model=m,
+                messages=[{"role":"system","content": _TXT_SYS},
+                          {"role":"user","content":
+                           f"Short idea: {user_prompt}\nSeed: {', '.join(seed_pos)}\nOutput: a single comma-separated tag list."}],
+                temperature=0.7, max_tokens=220
+            ) for m in LLM_MULTI_CANDIDATES]
+            rs = await asyncio.gather(*tasks, return_exceptions=True)
+            props = []
+            for r in rs:
+                try: props.append(r.choices[0].message.content)
+                except Exception: pass
+
+            if not props:
+                pos = ", ".join([user_prompt.strip()] + seed_pos)
+            else:
+                joined = "\n---\n".join(props)
+                merged = await client.chat.completions.create(
+                    model=self.aggregator,
+                    messages=[{"role":"system","content":
+                               "Merge candidate tag lists into ONE comma list (16–26 tags). Subject first; at most TWO style tags; keep high detailed/sharp focus/8k/UHD."},
+                              {"role":"user","content": joined}],
+                    temperature=0.2, max_tokens=240
+                )
+                raw = merged.choices[0].message.content
+                tags = _dedup_keep_order(_split_tags(raw))
+                subject = user_prompt.strip().rstrip(",.")
+                if subject and not any(subject.lower() == t.lower() for t in tags):
+                    tags = [subject] + tags
+                ordered = _order_tags([tags[0]], tags[1:])
+                pos = ", ".join(_dedup_keep_order(ordered))
+
+            # quality floor
+            for q in ["high detailed","sharp focus","8k","UHD"]:
+                if q.lower() not in {t.lower() for t in _split_tags(pos)}:
+                    pos += ", " + q
+
+            pos = _cleanup_commas(pos)
+            if clip77 and _count_tokens(pos) > 77:
+                pos = clip77_strict(pos, 77)
+
+            neg = ", ".join(seed_neg)
+            return pos, neg
+        finally:
+            await _maybe_close_async_together(client)
+
+    async def vlm_refine(self, image: Image.Image, original_prompt: str, components: List[str]) -> Dict[str, object]:
+        client = AsyncTogether(api_key=os.environ["TOGETHER_API_KEY"])
+        b64 = pil_to_base64(image, "PNG")
+
+        def _user_prompt_text() -> str:
+            return (
+                "You are a precise image-grounded critic.\n"
+                "1) List concrete visual problems and brief corrections.\n"
+                "2) Provide a refined prompt that keeps the original intent.\n\n"
+                f'Original prompt: "{original_prompt}"\n'
+                f"Key components to check: {components}\n"
+                "Output EXACTLY two tags:\n"
+                "<issues>...</issues>\n<refined>...</refined>"
+            )
+        try:
+            tasks = []
+            for m in VLM_CANDIDATES:
+                msgs = [
+                    {"role":"system","content": "Return ONLY <issues> and <refined>. No extra text."},
+                    {"role":"user","content": [
+                        {"type":"text","text": _user_prompt_text()},
+                        {"type":"image_url","image_url":{"url": f"data:image/png;base64,{b64}"}}
+                    ]}
+                ]
+                tasks.append(client.chat.completions.create(model=m, messages=msgs, temperature=0.2, max_tokens=420))
+
+            rs = await asyncio.gather(*tasks, return_exceptions=True)
+            ok = []
+            for m, r in zip(VLM_CANDIDATES, rs):
+                try: ok.append((m, r.choices[0].message.content))
+                except Exception: pass
+
+            if not ok:
+                return {"refined": original_prompt, "issues_merged": ""}
+
+            refined_items, per_vlm_issues = [], {}
+            for m, raw in ok:
+                issues = _extract_tag(raw, "issues", "")
+                refined = _extract_tag(raw, "refined", original_prompt)
+                if refined.strip(): refined_items.append((m, refined.strip()))
+                if issues.strip():  per_vlm_issues[m] = _summarize_issues_lines(issues, 5)
+
+            joined_issues  = "\n".join(f"[{m}] {t}" for m,t in per_vlm_issues.items())
+            joined_refined = "\n".join(f"[{m}] {t}" for m,t in refined_items) if refined_items else original_prompt
+
+            merged = await client.chat.completions.create(
+                model=self.aggregator,
+                messages=[{"role":"system","content":
+                           "Merge multiple critics. Output ONLY <issues> (≤5 bullets) and <refined> (≤70 words)."},
+                          {"role":"user","content": f"{joined_issues}\n\n----\n\n{joined_refined}"}],
+                temperature=0.2, max_tokens=420
+            )
+            final_raw = merged.choices[0].message.content
+            final_refined = clip77_strict(_extract_tag(final_raw, "refined", original_prompt), 77)
+            issues_merged = _summarize_issues_lines(_extract_tag(final_raw, "issues", ""), 5)
+
+            return {"refined": final_refined, "issues_merged": issues_merged}
+        finally:
+            await _maybe_close_async_together(client)
+
+    @staticmethod
+    def merge_vlm_multi_text(vlm_refined_77: str, tags_77: str) -> str:
+        vlm_tags = _split_tags(vlm_refined_77)
+        moa_tags = _split_tags(tags_77)
+        merged = _dedup_keep_order(_order_tags([vlm_tags[0] if vlm_tags else ""], (vlm_tags[1:] + moa_tags)))
+        merged = [t for t in merged if t]
+        text = _cleanup_commas(", ".join(merged))
+        if _count_tokens(text) > 77:
+            text = clip77_strict(text, 77)
+        return text
+
+# =========================
+# 5) Standalone frequency_fusion (your reference)
+# =========================
+@torch.no_grad()
+def frequency_fusion(
+    x_hi_latent: torch.Tensor,
+    x_lo_latent: torch.Tensor,
+    base_c: float = 0.5,
+    rho_t: float = 0.85,
+    device=None,
+) -> torch.Tensor:
+    if device is None:
+        device = x_hi_latent.device
+    B, C, H, W = x_hi_latent.shape
+
+    x_h = x_hi_latent.to(torch.float32).to(device)
+    x_l = x_lo_latent.to(torch.float32).to(device)
+
+    Xh = torch.fft.fftshift(torch.fft.fftn(x_h, dim=(-2, -1)), dim=(-2, -1))
+    Xl = torch.fft.fftshift(torch.fft.fftn(x_l, dim=(-2, -1)), dim=(-2, -1))
+
+    tau_h = int(H * base_c * (1 - rho_t))
+    tau_w = int(W * base_c * (1 - rho_t))
+
+    mask = torch.ones((B, C, H, W), device=device, dtype=torch.float32)
+    cy, cx = H // 2, W // 2
+    if tau_h > 0 and tau_w > 0:
+        mask[..., cy - tau_h : cy + tau_h, cx - tau_w : cx + tau_w] = rho_t
+
+    Xf = Xh * mask + Xl * (1 - mask)
+    x = torch.fft.ifftn(torch.fft.ifftshift(Xf, dim=(-2, -1)), dim=(-2, -1)).real
+    x = x + torch.randn_like(x) * 0.001
+    return x.to(dtype=x_hi_latent.dtype)
+
+def _decode_to_pil(latents, pipe):
+    out = decode_image_sdxl(latents, pipe)
+    if isinstance(out, Image.Image):
+        return out
+    if hasattr(out, "images"):
+        return out.images[0]
+    return out
+
+# =========================
+# 6) Helpers for variants
+# =========================
+def _normalize_enabled(enabled_variants: Optional[Iterable[str]]) -> Set[str]:
+    default = set(VARIANT_LABELS.keys())
+    if enabled_variants is None:
+        return default
+    return set(enabled_variants)
+
+def _guidance_for_k(k: int) -> float:
+    if k >= 20: return 12.0
+    if k >= 10: return 7.5
+    return 5.2
+
+async def _shared_materials(user_prompt: str, seed: int, H: int, W: int, preset: str):
+    critic = CritiCore(preset=preset)
+
+    pos_tags_77, neg_tags = await critic.make_tags(user_prompt, clip77=True)
+    comps = await critic.decompose_components(user_prompt)
+
+    z0_og,  base_og  = base_sample_latent(user_prompt, seed=seed, H=H, W=W, neg=DEFAULT_NEG)
+    z0_enh, base_enh = base_sample_latent(pos_tags_77, seed=seed, H=H, W=W, neg=neg_tags)
+
+    vlm_out = await critic.vlm_refine(base_enh, pos_tags_77, comps or [])
+    vlm_agg_77 = vlm_out.get("refined") or pos_tags_77
+
+    return dict(
+        pos_tags_77=pos_tags_77, neg_tags=neg_tags, comps=comps,
+        z0_og=z0_og, base_og=base_og,
+        z0_enh=z0_enh, base_enh=base_enh,
+        vlm_agg_77=vlm_agg_77,
+        critic=critic
+    )
+
+async def _collect_meta(user_prompt: str, seed: int, H: int, W: int, preset: str):
+    shared = await _shared_materials(user_prompt, seed, H, W, preset)
+    return {
+        "user_prompt": user_prompt,
+        "pos_tags_77": shared["pos_tags_77"],
+        "neg_tags": shared["neg_tags"],
+        "components": shared["comps"],
+        "vlm_agg_77_on_multi_llm": shared["vlm_agg_77"],
+    }
+
+# =========================
+# 7) Variants generator (signature matches your reference: last_k_list + guidance_list)
+# =========================
+async def generate_variants(
+    user_prompt: str,
+    seed: int,
+    H: int, W: int,
+    total_steps_refine: int,
+    last_k_list: Iterable[int],
+    guidance_list: Optional[List[float]] = None,
+    preset: str = "hq_preference",
+    out_dir: Optional[Path] = None,
+    enabled_variants: Optional[Iterable[str]] = None,
+) -> Dict[str, Dict[int, Image.Image]]:
+    enabled = _normalize_enabled(enabled_variants)
+
+    lk = int(last_k_list) if isinstance(last_k_list, int) else (int(list(last_k_list)[-1]) if last_k_list else 36)
+
+    shared = await _shared_materials(user_prompt, seed, H, W, preset)
+    pos_tags_77 = shared["pos_tags_77"]; comps = shared["comps"]
+    z0_og, base_og   = shared["z0_og"], shared["base_og"]
+    z0_enh, base_enh = shared["z0_enh"], shared["base_enh"]
+    vlm_agg_77       = shared["vlm_agg_77"]
+    critic: CritiCore = shared["critic"]
+
+    out: Dict[str, Dict[int, Image.Image]] = {}
+
+    def _save(im: Image.Image, vname: str, k: int = 0):
+        if out_dir is None: return
+        sub = out_dir / f"var_{vname}"
+        sub.mkdir(parents=True, exist_ok=True)
+        im.save(sub / f"{vname}_k{k}.png")
+
+    # 1) base_multi_llm
+    if "base_multi_llm" in enabled:
+        v = VARIANT_LABELS["base_multi_llm"]
+        out[v] = {0: base_enh}
+        _save(base_enh, v, 0)
+
+    # 2) criticore_on_multi_llm__specfusion
+    if "criticore_on_multi_llm__specfusion" in enabled:
+        v = VARIANT_LABELS["criticore_on_multi_llm__specfusion"]; out[v] = {}
+        refined_on_enh = CritiCore.merge_vlm_multi_text(vlm_agg_77, pos_tags_77)
+
+        strength = float(strength_for_last_k(lk, total_steps_refine))
+        guidance = float(guidance_list[-1]) if guidance_list else float(_guidance_for_k(lk))
+        steps    = int(total_steps_refine)
+
+        z_ref = img2img_latent(
+            refined_on_enh, z0_enh,
+            strength=strength, guidance=guidance, steps=steps,
+            seed=seed + 2100 + lk
+        )
+        fused_lat = frequency_fusion(z_ref, z0_enh, base_c=0.5, rho_t=RHO_T_DEFAULT, device=DEVICE)
+        img_sf = _decode_to_pil(fused_lat, SDXL_i2i)
+
+        out[v][0] = img_sf
+        _save(img_sf, v, 0)
+
+    # 3) base_original
+    if "base_original" in enabled:
+        v = VARIANT_LABELS["base_original"]
+        out[v] = {0: base_og}
+        _save(base_og, v, 0)
+
+    # 4) criticore_on_original__specfusion
+    if "criticore_on_original__specfusion" in enabled:
+        v = VARIANT_LABELS["criticore_on_original__specfusion"]; out[v] = {}
+
+        vlm_on_og = await critic.vlm_refine(base_og, user_prompt, comps or [])
+        refined_og_77 = clip77_strict(vlm_on_og.get("refined") or user_prompt, 77)
+        refined_merge = CritiCore.merge_vlm_multi_text(refined_og_77, pos_tags_77)
+
+        strength = float(strength_for_last_k(lk, total_steps_refine))
+        guidance = float(guidance_list[-1]) if guidance_list else float(_guidance_for_k(lk))
+        steps    = int(total_steps_refine)
+
+        z_ref = img2img_latent(
+            refined_merge, z0_og,
+            strength=strength, guidance=guidance, steps=steps,
+            seed=seed + 2400 + lk
+        )
+        fused_lat = frequency_fusion(z_ref, z0_og, base_c=0.5, rho_t=RHO_T_DEFAULT, device=DEVICE)
+        img_sf = _decode_to_pil(fused_lat, SDXL_i2i)
+
+        out[v][0] = img_sf
+        _save(img_sf, v, 0)
+
+    return out
+
+# =========================
+# 8) Full CADR pipeline (kept for your "完整 demo" tab)
+# =========================
+def _try_pref_score(enhanced_prompt: str, base_img: Image.Image) -> Optional[float]:
+    fn = globals().get("pref_score", None)  # optional external
+    if fn is None:
+        return None
+    try:
+        s01 = float(fn(enhanced_prompt, base_img))
+        return max(0.0, min(100.0, s01 * 100.0))
+    except Exception:
+        return None
+
+def _clamp01(x: float) -> float: return max(0.0, min(1.0, float(x)))
+def _lerp(a: float, b: float, t: float) -> float: return a + (b - a) * t
+
+class SpecFusionCADR:
+    def __init__(self, device: torch.device):
+        self.device = device
+
+    @staticmethod
+    def cadr_from_alignment(align_score: float) -> Tuple[float, float, int, float]:
+        s = _clamp01(align_score / 100.0); mis = 1.0 - s
+        strength = _lerp(0.12, 0.30, mis)
+        guidance = _lerp(3.6,  5.0,  mis)
+        steps    = int(round(_lerp(16, 30, mis)))
+        rho_t    = _lerp(0.60, 0.85, mis)
+        return strength, guidance, steps, rho_t
+
+    @torch.no_grad()
+    def final_touch(self, enhanced_prompt: str, base_latent: torch.Tensor, align_score: float, seed: int):
+        strength, guidance, steps, rho_t = self.cadr_from_alignment(float(align_score))
+        z_ref = img2img_latent(enhanced_prompt, base_latent, strength=strength, guidance=guidance, steps=steps, seed=seed)
+        fused = frequency_fusion(z_ref, base_latent, base_c=0.5, rho_t=rho_t, device=DEVICE)
+        img = decode_image_sdxl(fused, SDXL_i2i)
+        return img, dict(strength=strength, guidance=guidance, steps=steps, rho_t=rho_t)
+
+async def pipeline_full_cadr(user_prompt: str, seed: int, H: int, W: int, preset: str, align_score: Optional[float], save_dir: Optional[Path]):
+    critic = CritiCore(preset=preset)
+    spec   = SpecFusionCADR(device=DEVICE)
+
+    z0_base, base_img = base_sample_latent(user_prompt, seed=seed, H=H, W=W, neg=DEFAULT_NEG)
+
+    pos_tags_77, neg_tags = await critic.make_tags(user_prompt, clip77=True)
+    comps = await critic.decompose_components(user_prompt)
+    vlm_out = await critic.vlm_refine(base_img, user_prompt, comps or [])
+    vlm_ref_77 = vlm_out.get("refined") or user_prompt
+    enhanced_77 = CritiCore.merge_vlm_multi_text(vlm_ref_77, pos_tags_77)
+
+    if align_score is None:
+        auto = _try_pref_score(enhanced_77, base_img)
+        align_score = auto if auto is not None else 60.0
+
+    final_img, cadr_params = spec.final_touch(enhanced_77, z0_base, float(align_score), seed=seed+999)
+
+    meta = {
+        "user_prompt": user_prompt,
+        "pos_tags_77": pos_tags_77,
+        "vlm_refined_77": vlm_ref_77,
+        "enhanced_prompt_77": enhanced_77,
+        "align_score": float(align_score),
+        "cadr_params": cadr_params,
+        "components": comps,
+        "vlm_issues": vlm_out.get("issues_merged",""),
+    }
+    meta_json = json.dumps(meta, ensure_ascii=False, indent=2)
+
+    if save_dir is not None:
+        save_dir.mkdir(parents=True, exist_ok=True)
+        base_img.save(save_dir / "base.png")
+        final_img.save(save_dir / "cadr_final.png")
+        (save_dir / "record.json").write_text(meta_json, encoding="utf-8")
+
+    return base_img, final_img, enhanced_77, meta_json
+
+# =========================
+# 9) UI callbacks (reference style)
+# =========================
+def ui_run_once(
+    user_prompt: str,
+    seed: int,
+    H: int,
+    W: int,
+    preset: str,
+    total_steps_refine: int,
+    last_k: int,
+    guidance: float,
+    enabled_variants: List[str],
+    save_outputs: bool,
+    out_dir: str,
+):
+    t0 = time.time()
+    try:
+        if not user_prompt or not user_prompt.strip():
+            return [], "Empty prompt."
+
+        if not TOGETHER_API_KEY:
+            return [], "ERROR: TOGETHER_API_KEY not set."
+
+        # display -> internal
+        display_to_internal = {v: k for k, v in VARIANT_LABELS.items()}
+        internal_enabled = [display_to_internal.get(v, v) for v in (enabled_variants or [])]
+
+        out_path = Path(out_dir) if (save_outputs and out_dir) else None
+        if out_path is not None:
+            out_path.mkdir(parents=True, exist_ok=True)
+
+        results = _run_async(generate_variants(
+            user_prompt=user_prompt.strip(),
+            seed=int(seed),
+            H=int(H), W=int(W),
+            total_steps_refine=int(total_steps_refine),
+            last_k_list=(int(last_k),),
+            guidance_list=[float(guidance)] if guidance > 0 else None,
+            preset=preset,
+            out_dir=out_path,
+            enabled_variants=internal_enabled,
+        ))
+
+        gallery = []
+        for name in VARIANT_ORDER:
+            if name in results and 0 in results[name]:
+                gallery.append((results[name][0], name))
+
+        try:
+            meta = _run_async(_collect_meta(user_prompt.strip(), int(seed), int(H), int(W), preset))
+        except Exception as e:
+            meta = {"meta_error": str(e)}
+
+        meta["ui"] = {
+            "seed": int(seed),
+            "H": int(H),
+            "W": int(W),
+            "preset": preset,
+            "total_steps_refine": int(total_steps_refine),
+            "last_k": int(last_k),
+            "guidance": float(guidance),
+            "enabled_variants": enabled_variants,
+            "save_outputs": bool(save_outputs),
+            "out_dir": out_dir if save_outputs else None,
+        }
+        meta["elapsed_sec"] = round(time.time() - t0, 3)
+
+        return gallery, json.dumps(meta, ensure_ascii=False, indent=2)
+
+    except Exception:
+        return [], traceback.format_exc()
+
+def ui_run_full(
+    user_prompt: str,
+    seed: int,
+    H: int,
+    W: int,
+    preset: str,
+    align_mode: str,
+    align_score: float,
+    save_outputs: bool,
+    out_dir: str,
+):
+    try:
+        if not user_prompt or not user_prompt.strip():
+            return None, None, "", "Empty prompt."
+        if not TOGETHER_API_KEY:
+            return None, None, "", "ERROR: TOGETHER_API_KEY not set."
+
+        save_dir = Path(out_dir) if (save_outputs and out_dir) else None
+        a = None if align_mode.startswith("Auto") else float(align_score)
+
+        base_img, final_img, enhanced_77, meta_json = _run_async(
+            pipeline_full_cadr(
+                user_prompt=user_prompt.strip(),
+                seed=int(seed), H=int(H), W=int(W),
+                preset=preset,
+                align_score=a,
+                save_dir=save_dir,
+            )
+        )
+        return base_img, final_img, enhanced_77, meta_json
+    except Exception:
+        return None, None, "", traceback.format_exc()
+
+# =========================
+# 10) Gradio UI (matches your reference fixes)
+# =========================
+with gr.Blocks(title="CritiFusion (SDXL) Demo", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+    "## CritiFusion Demo (SDXL)\n"
+    "- You can run **Variants** (outputs 4 variants in one click).\n"
+    "- You can also run the **Full CADR Pipeline** (end-to-end).\n"
+    f"- Device: **{DEVICE_STR}**, DType: **{DTYPE}**\n"
+    f"- Together API: {'✅ set' if TOGETHER_API_KEY else '❌ missing (set TOGETHER_API_KEY)'}"
+    )
+
+    with gr.Tabs():
+        with gr.Tab("Variants (Run Once)"):
+            with gr.Row():
+                with gr.Column(scale=7):
+                    user_prompt = gr.Textbox(
+                        label="Prompt",
+                        value="A fluffy orange cat lying on a window ledge, front-facing, stylized in 3D Pixar look, soft indoor lighting",
+                        lines=3,
+                    )
+                    with gr.Row():
+                        seed = gr.Number(label="Seed", value=2026, precision=0)
+                        preset = gr.Dropdown(label="Preset", choices=["hq_preference"], value="hq_preference")
+                    with gr.Row():
+                        H = gr.Number(label="H", value=1024, precision=0)
+                        W = gr.Number(label="W", value=1024, precision=0)
+                    with gr.Row():
+                        total_steps_refine = gr.Slider(label="total_steps_refine", minimum=10, maximum=80, step=1, value=50)
+                        last_k = gr.Slider(label="last_k", minimum=1, maximum=50, step=1, value=37)
+                    guidance = gr.Slider(label="Guidance (0 => fallback rule)", minimum=0.0, maximum=15.0, step=0.1, value=0.0)
+
+                    enabled_variants = gr.CheckboxGroup(
+                        label="Enabled Variants",
+                        choices=[VARIANT_LABELS[k] for k in VARIANT_LABELS.keys()],
+                        value=[
+                            VARIANT_LABELS["base_original"],
+                            VARIANT_LABELS["base_multi_llm"],
+                            VARIANT_LABELS["criticore_on_original__specfusion"],
+                            VARIANT_LABELS["criticore_on_multi_llm__specfusion"],
+                        ],
+                    )
+
+                    with gr.Row():
+                        save_outputs = gr.Checkbox(label="Save outputs to disk", value=False)
+                        out_dir = gr.Textbox(label="Output dir (only if save enabled)", value="./variants_demo_gradio")
+
+                    run_btn = gr.Button("Run", variant="primary")
+
+                with gr.Column(scale=8):
+                    gallery = gr.Gallery(label="Results", columns=2, height=600)
+                    meta_json = gr.Code(label="Meta / Debug (JSON)", language="json")
+
+            run_btn.click(
+                fn=ui_run_once,
+                inputs=[user_prompt, seed, H, W, preset, total_steps_refine, last_k, guidance, enabled_variants, save_outputs, out_dir],
+                outputs=[gallery, meta_json],
+                api_name=False,
+                show_api=False,
+            )
+
+        with gr.Tab("Full CADR Pipeline"):
+            with gr.Row():
+                with gr.Column(scale=7):
+                    p2 = gr.Textbox(
+                        label="Prompt",
+                        value="A fluffy orange cat lying on a window ledge, front-facing, stylized in 3D Pixar look, soft indoor lighting",
+                        lines=3,
+                    )
+                    with gr.Row():
+                        seed2 = gr.Number(label="Seed", value=2026, precision=0)
+                        preset2 = gr.Dropdown(label="Preset", choices=["hq_preference"], value="hq_preference")
+                    with gr.Row():
+                        H2 = gr.Number(label="H", value=1024, precision=0)
+                        W2 = gr.Number(label="W", value=1024, precision=0)
+
+                    align_mode = gr.Radio(
+                        label="Alignment score source",
+                        choices=["Auto (pref_score if available else 60)", "Manual (use slider below)"],
+                        value="Auto (pref_score if available else 60)",
+                    )
+                    align_score = gr.Slider(label="Manual align_score (0..100)", minimum=0, maximum=100, step=1, value=60)
+
+                    with gr.Row():
+                        save2 = gr.Checkbox(label="Save outputs to disk", value=False)
+                        out2 = gr.Textbox(label="Output dir (only if save enabled)", value="./full_cadr_gradio")
+
+                    run2 = gr.Button("Run Full CADR", variant="primary")
+
+                with gr.Column(scale=8):
+                    base_img = gr.Image(label="Base", type="pil")
+                    final_img = gr.Image(label="Final (CADR + SpecFusion)", type="pil")
+                    enhanced = gr.Textbox(label="Enhanced prompt (≤77)", lines=3)
+                    meta2 = gr.Code(label="Meta / Debug (JSON)", language="json")
+
+            run2.click(
+                fn=ui_run_full,
+                inputs=[p2, seed2, H2, W2, preset2, align_mode, align_score, save2, out2],
+                outputs=[base_img, final_img, enhanced, meta2],
+                api_name=False,
+                show_api=False,
+            )
+
+# IMPORTANT: share=True fixes "localhost not accessible"
+demo.queue().launch(debug=True, share=True, max_threads=1, show_api=False)

requirements.txt

@@ -0,0 +1,11 @@
+gradio>=4.44.0
+torch
+diffusers>=0.30.0
+transformers>=4.43.0
+accelerate
+safetensors
+huggingface_hub
+pillow
+numpy
+together
+nest_asyncio