feat: switch both pipelines to custom dataset checkpoints

Load single-file checkpoints from the nsfwalex/checkpoint_n_lora dataset
instead of the stock HF models:
- Z-Image pipeline -> Moody Pro Mix (ZIT V12 DPO), a transformer-only
Z-Image-Turbo finetune swapped into the base pipeline via
ZImageTransformer2DModel.from_single_file.
- SDXL pipeline -> One Obsession v2.1 Anime (Illustrious) via
StableDiffusionXLPipeline.from_single_file, keeping the fp16-fix VAE
and Compel prompt weighting.

Relabels the model dropdown and helper text to the new model names.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -1,4 +1,5 @@
 import io
+import os
 import random
 import re
 
@@ -6,8 +7,14 @@ import numpy as np
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline, AutoPipelineForText2Image, AutoencoderKL
+from diffusers import (
+    DiffusionPipeline,
+    StableDiffusionXLPipeline,
+    AutoencoderKL,
+    ZImageTransformer2DModel,
+)
 from compel import Compel, ReturnedEmbeddingsType
+from huggingface_hub import hf_hub_download
 
 import r2_uploader
 
@@ -19,13 +26,39 @@ R2_NAMESPACE = "s01"
 # =============================================================================
 # Model registry
 # =============================================================================
-# Two pipelines are exposed through a single UI via the selector below:
-#   * Z-Image-Turbo  -> fast, guidance-free distilled model (8 DiT forwards)
-#   * NoobAI-XL v1.1 -> SDXL anime model with Compel prompt weighting
-MODEL_ZIMAGE = "Z-Image-Turbo"
-MODEL_NOOBXL = "NoobAI-XL v1.1"
+# Two pipelines are exposed through a single UI via the selector below. Both now
+# run custom single-file checkpoints pulled from the nsfwalex/checkpoint_n_lora
+# dataset (populated by the model_downloader Space):
+#   * Moody Pro Mix     -> Z-Image-Turbo finetune (fast, guidance-free DiT)
+#   * One Obsession     -> Illustrious / SDXL anime model with Compel weighting
+# Constant names keep their pipeline meaning (zimage = the Z-Image DiT pipeline,
+# noobxl = the SDXL pipeline); only the displayed labels / weights changed.
+MODEL_ZIMAGE = "Moody Pro Mix (ZIT V12 DPO)"
+MODEL_NOOBXL = "One Obsession v2.1 Anime"
 MODEL_CHOICES = [MODEL_ZIMAGE, MODEL_NOOBXL]
 
+# Custom checkpoints live in this HF dataset; HF_TOKEN (a Space secret) grants
+# read access. Paths mirror the model_downloader index.json layout.
+CHECKPOINT_DATASET = "nsfwalex/checkpoint_n_lora"
+ZIMAGE_CHECKPOINT = "checkpoints/zimageturbo/moodyProMix_zitV12DPO.safetensors"
+NOOBXL_CHECKPOINT = "checkpoints/il/oneObsession_v21Anime.safetensors"
+# Base Z-Image repo supplies the VAE / text encoder / tokenizer / scheduler and
+# the transformer config; the single-file checkpoint only carries the DiT.
+ZIMAGE_BASE = "Tongyi-MAI/Z-Image-Turbo"
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+
+def _checkpoint_path(filename):
+    """Download a checkpoint from the dataset and return its local cache path."""
+    return hf_hub_download(
+        repo_id=CHECKPOINT_DATASET,
+        repo_type="dataset",
+        filename=filename,
+        token=HF_TOKEN,
+    )
+
+
 MAX_SEED = np.iinfo(np.int32).max
 
 NOOBXL_NEGATIVE = (
@@ -38,20 +71,32 @@ NOOBXL_NEGATIVE = (
 # Load both pipelines once at startup. On ZeroGPU the `.to("cuda")` calls are
 # captured by the runtime and the work happens inside the @spaces.GPU function.
 # -----------------------------------------------------------------------------
-print("Loading Z-Image-Turbo pipeline...")
+print("Loading Z-Image pipeline (Moody Pro Mix)...")
+# The CivitAI checkpoint is a transformer-only (model.diffusion_model.*) file, so
+# load the base Z-Image-Turbo pipeline for the other components and swap in the
+# custom DiT loaded from the single-file checkpoint.
+zimage_transformer = ZImageTransformer2DModel.from_single_file(
+    _checkpoint_path(ZIMAGE_CHECKPOINT),
+    config=ZIMAGE_BASE,
+    subfolder="transformer",
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN,
+)
 zimage_pipe = DiffusionPipeline.from_pretrained(
-    "Tongyi-MAI/Z-Image-Turbo",
+    ZIMAGE_BASE,
+    transformer=zimage_transformer,
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
 zimage_pipe.to("cuda")
 
-print("Loading NoobAI-XL v1.1 pipeline...")
+print("Loading SDXL pipeline (One Obsession)...")
+# Override the in-checkpoint VAE with the fp16-fix VAE to avoid black/NaN images.
 noobxl_vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 )
-noobxl_pipe = AutoPipelineForText2Image.from_pretrained(
-    "Menyu/noobaiXLNAIXL_epsilonPred11Version",
+noobxl_pipe = StableDiffusionXLPipeline.from_single_file(
+    _checkpoint_path(NOOBXL_CHECKPOINT),
     vae=noobxl_vae,
     torch_dtype=torch.float16,
     use_safetensors=True,
@@ -397,7 +442,7 @@ def apply_model_defaults(model_name):
     return (
         gr.update(value=d["steps"]),
         gr.update(value=d["guidance"], interactive=is_noob,
-                  info="Z-Image-Turbo is guidance-free (forced to 0)" if not is_noob else "Classifier-free guidance"),
+                  info="Moody Pro Mix (Z-Image) is guidance-free (forced to 0)" if not is_noob else "Classifier-free guidance"),
         gr.update(value=d["height"]),
         gr.update(value=d["width"]),
         gr.update(visible=is_noob),
@@ -484,7 +529,7 @@ with gr.Blocks(fill_height=True) as demo:
                 choices=MODEL_CHOICES,
                 value=MODEL_ZIMAGE,
                 label="🧠 Model",
-                info="Z-Image-Turbo: fast & general • NoobAI-XL: anime / illustration",
+                info="Moody Pro Mix: fast & general (Z-Image) • One Obsession: anime / illustration (SDXL)",
             )
 
             prompt = gr.Textbox(
@@ -522,13 +567,13 @@ with gr.Blocks(fill_height=True) as demo:
                 num_inference_steps = gr.Slider(
                     minimum=1, maximum=50, value=9, step=1,
                     label="Inference Steps",
-                    info="Z-Image: ~9 • NoobAI-XL: ~28",
+                    info="Moody Pro Mix: ~9 • One Obsession: ~28",
                 )
 
                 guidance_scale = gr.Slider(
                     minimum=0.0, maximum=10.0, value=0.0, step=0.1,
                     label="Guidance Scale",
-                    info="Z-Image-Turbo is guidance-free (forced to 0)",
+                    info="Moody Pro Mix (Z-Image) is guidance-free (forced to 0)",
                     interactive=False,
                 )
 
@@ -567,8 +612,8 @@ with gr.Blocks(fill_height=True) as demo:
         ---
         <div style="text-align: center; opacity: 0.7; font-size: 0.9em; margin-top: 1rem;">
         <strong>Models:</strong>
-        <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank">Z-Image-Turbo</a> •
-        <a href="https://huggingface.co/Menyu/noobaiXLNAIXL_epsilonPred11Version" target="_blank">NoobAI-XL v1.1</a> •
+        Moody Pro Mix (ZIT V12 DPO, Z-Image-Turbo finetune) •
+        One Obsession v2.1 Anime (Illustrious / SDXL) •
         <strong>Demo by:</strong> <a href="https://x.com/realmrfakename" target="_blank">@mrfakename</a>
         </div>
         """,