Side-by-side: Nanosaur2 BF16/INT8 + Supra2-IMG

README.md

@@ -1,5 +1,5 @@
 ---
-title: Nanosaur2 670M
+title: Small T2I Models Side by Side
 emoji: 🦖
 colorFrom: green
 colorTo: yellow
@@ -9,17 +9,21 @@ app_file: app.py
 pinned: false
 license: mit
 hardware: zero-a10g
-short_description: Tiny 670M anime/furry text-to-image model (ComfyUI nodes)
+short_description: Compare Nanosaur2 (BF16/INT8) and Supra2-IMG
 python_version: "3.12"
 startup_duration_timeout: "30m"
 models:
   - well9472/Nanosaur2-670M
+  - bertbobson/Nanosaur2-670M-INT8-ConvRot
+  - SupraLabs/Supra2-IMG
 ---
 
-# Nanosaur2 670M
+# Small text-to-image models, side by side
 
-Demo for [well9472/Nanosaur2-670M](https://huggingface.co/well9472/Nanosaur2-670M): a 670M-parameter illustration text-to-image DiT (Gemma3-270M text encoder, DINOv2 semantic VAE), trained on e621/Danbooru. Tags or natural language.
+Runs the same prompt and seed through new tiny text-to-image models:
 
-The model only ships as ComfyUI custom nodes, so this Space runs ComfyUI as a library (cloned at startup) with the author's `nanosaur2_support` nodes. No weights are stored in this repo: everything is downloaded from the Hub.
+- [well9472/Nanosaur2-670M](https://huggingface.co/well9472/Nanosaur2-670M): 670M illustration DiT (ComfyUI nodes, run here through ComfyUI as a library)
+- [bertbobson/Nanosaur2-670M-INT8-ConvRot](https://huggingface.co/bertbobson/Nanosaur2-670M-INT8-ConvRot): INT8 version of the same model (patched nodes.py)
+- [SupraLabs/Supra2-IMG](https://huggingface.co/SupraLabs/Supra2-IMG): 100M DiT, 256x256
 
-Defaults follow the model's own workflow: Euler / simple, CFG 4, 30 steps, "alternate" guidance (CFG + path-drop).
+No weights are stored in this repo: ComfyUI is cloned and all models are downloaded from the Hub at startup.

app.py

@@ -1,22 +1,36 @@
+import importlib.util
 import os
 import random
+import shutil
 import subprocess
 import sys
 
 import gradio as gr
 import spaces
-import torch  # noqa: F401  (must be imported before CUDA is touched, as required by ZeroGPU)
-from huggingface_hub import snapshot_download
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image
 
 # --------------------------------------------------------------------------------------
-# Nanosaur2-670M ships only as ComfyUI custom nodes (nanosaur2_support). To run the exact
-# same code path, this Space uses ComfyUI as a library. Nothing is stored in this repo:
-# ComfyUI is cloned and the model files are downloaded from the Hub at startup.
+# Side-by-side comparison of small, new text-to-image models:
+#   * Nanosaur2-670M (BF16)         - ComfyUI custom nodes (author's code)
+#   * Nanosaur2-670M INT8 ConvRot   - same model, INT8 quant (needs patched nodes.py)
+#   * Supra2-IMG (~100M)            - standalone DiT, 256x256
+# Nothing is stored in this repo: ComfyUI is cloned and all weights come from the Hub.
 # --------------------------------------------------------------------------------------
-MODEL_REPO = "well9472/Nanosaur2-670M"
-COMFY_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ComfyUI")
+NANO_REPO = "well9472/Nanosaur2-670M"
+INT8_REPO = "bertbobson/Nanosaur2-670M-INT8-ConvRot"
+SUPRA_REPO = "SupraLabs/Supra2-IMG"
+HERE = os.path.dirname(os.path.abspath(__file__))
+COMFY_DIR = os.path.join(HERE, "ComfyUI")
+PKG_DIR = os.path.join(HERE, "pkgs")
 MAX_SEED = 2**31 - 1
 
+NANO_BF16 = "Nanosaur2 670M (BF16)"
+NANO_INT8 = "Nanosaur2 670M (INT8 ConvRot)"
+SUPRA = "Supra2-IMG (100M, 256px)"
+ALL_MODELS = [NANO_BF16, NANO_INT8, SUPRA]
+
 DEFAULT_PROMPT = (
     "newest, masterpiece, 1girl, solo, (fennec ears:1.3), long blonde wavy hair, blue eyes, "
     "big fluffy tail, smile, forest, sunlight"
@@ -26,130 +40,215 @@ DEFAULT_NEGATIVE = (
     "text, bad anatomy, deformed, extra limbs, missing fingers, cropped"
 )
 
+# ---------------------------------- startup (no CUDA) ----------------------------------
 if not os.path.isdir(COMFY_DIR):
     subprocess.run(
         ["git", "clone", "--depth", "1", "https://github.com/comfyanonymous/ComfyUI", COMFY_DIR],
         check=True,
     )
 
-MODEL_DIR = snapshot_download(
-    MODEL_REPO,
-    allow_patterns=["*.safetensors", "nanosaur2_support/*.py"],
+NANO_DIR = snapshot_download(NANO_REPO, allow_patterns=["*.safetensors", "nanosaur2_support/*.py"])
+INT8_DIR = snapshot_download(INT8_REPO, allow_patterns=["nanosaur2_int8_conv.safetensors", "nodes.py"])
+
+# Two copies of the author's node package: the original, and one with the INT8-patched nodes.py.
+os.makedirs(PKG_DIR, exist_ok=True)
+for pkg in ("nanosaur2_support", "nanosaur2_support_int8"):
+    dst = os.path.join(PKG_DIR, pkg)
+    if not os.path.isdir(dst):
+        shutil.copytree(os.path.join(NANO_DIR, "nanosaur2_support"), dst, symlinks=False)
+shutil.copyfile(
+    os.path.join(INT8_DIR, "nodes.py"),
+    os.path.join(PKG_DIR, "nanosaur2_support_int8", "nodes.py"),
 )
 
-_runtime = {}
+SUPRA_CODE = hf_hub_download(SUPRA_REPO, "inference.py")
+SUPRA_CKPT = hf_hub_download(SUPRA_REPO, "model_final_ema.pt")
+
+_rt = {}
 
 
-def get_runtime():
-    """Import ComfyUI lazily (inside the GPU context) and register the Nanosaur2 nodes."""
-    if _runtime:
-        return _runtime
+# ------------------------------------- Nanosaur2 ---------------------------------------
+def get_comfy():
+    if "nodes" in _rt:
+        return _rt
     sys.argv = [sys.argv[0]]  # ComfyUI parses sys.argv on import
-    for p in (COMFY_DIR, MODEL_DIR):
+    for p in (COMFY_DIR, PKG_DIR):
         if p not in sys.path:
             sys.path.insert(0, p)
-
     import folder_paths
     import nodes
 
     for folder in ("diffusion_models", "text_encoders", "vae"):
-        folder_paths.add_model_folder_path(folder, MODEL_DIR)
+        folder_paths.add_model_folder_path(folder, NANO_DIR)
+    folder_paths.add_model_folder_path("diffusion_models", INT8_DIR)
 
     from nanosaur2_support.nodes import Nanosaur2Loader
+    from nanosaur2_support_int8.nodes import Nanosaur2Loader as Nanosaur2LoaderInt8
 
-    _runtime.update(nodes=nodes, Loader=Nanosaur2Loader)
-    return _runtime
+    _rt.update(nodes=nodes, Loader=Nanosaur2Loader, LoaderInt8=Nanosaur2LoaderInt8)
+    return _rt
 
 
 def snap(v, m=16):
     return max(m, int(round(v / m)) * m)
 
 
-@spaces.GPU(duration=90)
-def generate(
-    prompt,
-    negative_prompt,
-    width,
-    height,
-    steps,
-    cfg,
-    guidance_mode,
-    seed,
-    randomize_seed,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if not prompt or not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-
-    rt = get_runtime()
+def run_nanosaur(int8, prompt, negative, width, height, steps, cfg, guidance_mode, seed):
+    rt = get_comfy()
     nodes = rt["nodes"]
-    from PIL import Image
-
+    loader = rt["LoaderInt8"] if int8 else rt["Loader"]
+    unet = "nanosaur2_int8_conv.safetensors" if int8 else "nanosaur2_diffusion_model.safetensors"
     with torch.inference_mode():
-        model, clip, vae = rt["Loader"]().load(
-            "nanosaur2_diffusion_model.safetensors",
-            "nanosaur2_text_encoder.safetensors",
-            "nanosaur2_vae.safetensors",
-            guidance_mode,
+        model, clip, vae = loader().load(
+            unet, "nanosaur2_text_encoder.safetensors", "nanosaur2_vae.safetensors", guidance_mode
         )
-        positive = nodes.CLIPTextEncode().encode(clip, prompt)[0]
-        negative = nodes.CLIPTextEncode().encode(clip, negative_prompt or "")[0]
+        pos = nodes.CLIPTextEncode().encode(clip, prompt)[0]
+        neg = nodes.CLIPTextEncode().encode(clip, negative or "")[0]
         latent = nodes.EmptyLatentImage().generate(snap(width), snap(height), 1)[0]
         samples = nodes.KSampler().sample(
-            model, seed, int(steps), float(cfg), "euler", "simple", positive, negative, latent, 1.0
+            model, seed, int(steps), float(cfg), "euler", "simple", pos, neg, latent, 1.0
         )[0]
         image = nodes.VAEDecode().decode(vae, samples)[0]
-
     arr = (image[0].clamp(0, 1).cpu().float().numpy() * 255).round().astype("uint8")
-    return Image.fromarray(arr), seed
-
-
-CSS = """
-#page { max-width: 1100px; margin: 0 auto; }
-footer { display: none !important; }
-"""
-
-with gr.Blocks(title="Nanosaur2 670M") as demo:
+    try:
+        import comfy.model_management as mm
+
+        mm.unload_all_models()
+        mm.soft_empty_cache()
+    except Exception:
+        pass
+    return Image.fromarray(arr)
+
+
+# -------------------------------------- Supra2 -----------------------------------------
+def run_supra(prompt, cfg, steps, seed):
+    from diffusers import AutoencoderKL
+    from transformers import AutoTokenizer, T5EncoderModel
+
+    spec = importlib.util.spec_from_file_location("supra_inference", SUPRA_CODE)
+    sup = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(sup)
+
+    dev = torch.device("cuda")
+    state = torch.load(SUPRA_CKPT, map_location=dev, weights_only=False)
+    scfg = state.get("config", {}) if isinstance(state, dict) else {}
+    weights = state["ema"] if isinstance(state, dict) and "ema" in state else state.get("model", state)
+    model = sup.SupraDiT().to(dev).eval()
+    model.load_state_dict(weights, strict=True)
+    ctx_len = int(scfg.get("ctx_len", sup.MAX_CTX_LEN))
+
+    tokenizer = AutoTokenizer.from_pretrained(sup.T5_NAME)
+    text_model = T5EncoderModel.from_pretrained(sup.T5_NAME).to(dev).eval()
+    vae = AutoencoderKL.from_pretrained(sup.VAE_NAME).to(dev).eval()
+
+    def encode(texts):
+        tok = tokenizer(texts, padding="max_length", truncation=True, max_length=ctx_len, return_tensors="pt").to(dev)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            ctx = text_model(**tok).last_hidden_state.float()
+        return ctx, tok["attention_mask"].float()
+
+    with torch.no_grad():
+        ctx, cmask = encode([prompt])
+        use_cfg = cfg > 1.0
+        if use_cfg:
+            if "uncond_text" in scfg:
+                uctx = scfg["uncond_text"].to(dev).float().unsqueeze(0)
+                umask = scfg["uncond_mask"].to(dev).float().unsqueeze(0)
+            else:
+                uctx, umask = encode([""])
+            ctx_all, mask_all = torch.cat([ctx, uctx], 0), torch.cat([cmask, umask], 0)
+        torch.manual_seed(seed)
+        z = torch.randn(1, sup.LATENT_CH, sup.LATENT_SIZE, sup.LATENT_SIZE, device=dev)
+        dt = 1.0 / steps
+        for i in range(int(steps)):
+            t = torch.full((1,), i * dt, device=dev)
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if use_cfg:
+                    v_both = model(torch.cat([z, z], 0), torch.cat([t, t], 0), ctx_all, mask_all)
+                    v_c, v_u = v_both.float().chunk(2, 0)
+                    v = v_u + cfg * (v_c - v_u)
+                else:
+                    v = model(z, t, ctx, cmask).float()
+            z = z + dt * v
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            img = vae.decode(z / sup.VAE_SCALE).sample
+        img = ((img.clamp(-1, 1) + 1) / 2)[0].permute(1, 2, 0).float().cpu().numpy()
+    return Image.fromarray((img * 255).round().astype("uint8"))
+
+
+# --------------------------------------- UI --------------------------------------------
+@spaces.GPU(duration=120)
+def generate(
+    models, prompt, negative_prompt, width, height, steps, cfg, guidance_mode,
+    supra_steps, supra_cfg, seed, randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    if not models:
+        raise gr.Error("Select at least one model.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    results = []
+    for name in ALL_MODELS:
+        if name not in models:
+            continue
+        try:
+            if name == SUPRA:
+                img = run_supra(prompt, float(supra_cfg), int(supra_steps), seed)
+            else:
+                img = run_nanosaur(name == NANO_INT8, prompt, negative_prompt, width, height,
+                                   steps, cfg, guidance_mode, seed)
+            results.append((img, name))
+        except Exception as e:  # keep the other models' results
+            gr.Warning(f"{name} failed: {type(e).__name__}: {e}")
+    if not results:
+        raise gr.Error("All selected models failed, see the logs.")
+    return results, seed
+
+
+CSS = "#page { max-width: 1200px; margin: 0 auto; } footer { display: none !important; }"
+
+with gr.Blocks(title="Small T2I models side by side") as demo:
     with gr.Column(elem_id="page"):
         gr.Markdown(
-            "# Nanosaur2 670M\n"
-            "Tiny (670M DiT + Gemma3-270M + DINOv2 VAE) illustration text-to-image model trained on "
-            "e621/Danbooru. Accepts tags or natural language. "
-            "Model: [well9472/Nanosaur2-670M](https://huggingface.co/well9472/Nanosaur2-670M) (MIT, research use). "
-            "Tip: start the positive prompt with *newest, masterpiece* and the negative with *oldest, low quality*; "
-            "upweight with *(tag:1.3)*."
+            "# Small text-to-image models, side by side\n"
+            "Same prompt and seed on new tiny models: "
+            "[Nanosaur2-670M](https://huggingface.co/well9472/Nanosaur2-670M) (anime/furry, BF16 and "
+            "[INT8](https://huggingface.co/bertbobson/Nanosaur2-670M-INT8-ConvRot)) and "
+            "[Supra2-IMG](https://huggingface.co/SupraLabs/Supra2-IMG) (100M, 256px). "
+            "Nanosaur2 likes tags: start with *newest, masterpiece*, negative with *oldest, low quality*."
         )
         with gr.Row(equal_height=False):
-            with gr.Column():
+            with gr.Column(scale=2):
+                models = gr.CheckboxGroup(ALL_MODELS, value=[NANO_BF16, NANO_INT8, SUPRA], label="Models")
                 prompt = gr.Textbox(label="Prompt", lines=3, value=DEFAULT_PROMPT)
-                negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value=DEFAULT_NEGATIVE)
+                negative_prompt = gr.Textbox(label="Negative prompt (Nanosaur2 only)", lines=2, value=DEFAULT_NEGATIVE)
                 run = gr.Button("Generate", variant="primary", size="lg")
-                with gr.Accordion("Advanced", open=False):
+                with gr.Accordion("Nanosaur2 settings", open=False):
                     with gr.Row():
                         width = gr.Slider(512, 1536, value=832, step=16, label="Width")
                         height = gr.Slider(512, 1536, value=1216, step=16, label="Height")
-                    steps = gr.Slider(1, 80, value=30, step=1, label="Steps", info="Model card: 30-50 (Euler / simple)")
+                    steps = gr.Slider(1, 80, value=30, step=1, label="Steps")
                     cfg = gr.Slider(1.0, 12.0, value=4.0, step=0.1, label="CFG scale")
-                    guidance_mode = gr.Radio(
-                        ["alternate", "cfg", "path_drop"],
-                        value="alternate",
-                        label="Guidance mode",
-                        info="alternate = CFG on even steps, path-drop on odd steps (recommended)",
-                    )
+                    guidance_mode = gr.Radio(["alternate", "cfg", "path_drop"], value="alternate", label="Guidance mode")
+                with gr.Accordion("Supra2-IMG settings", open=False):
+                    supra_steps = gr.Slider(1, 100, value=50, step=1, label="Steps")
+                    supra_cfg = gr.Slider(1.0, 10.0, value=3.0, step=0.1, label="CFG scale")
+                with gr.Accordion("Seed", open=False):
                     seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                     randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
-            with gr.Column():
-                result = gr.Image(label="Result", format="png")
+            with gr.Column(scale=3):
+                gallery = gr.Gallery(label="Results", columns=3, height=620, object_fit="contain", format="png")
                 used_seed = gr.Number(label="Seed used", interactive=False)
 
         gr.on(
             [run.click, prompt.submit],
             generate,
-            [prompt, negative_prompt, width, height, steps, cfg, guidance_mode, seed, randomize_seed],
-            [result, used_seed],
+            [models, prompt, negative_prompt, width, height, steps, cfg, guidance_mode,
+             supra_steps, supra_cfg, seed, randomize_seed],
+            [gallery, used_seed],
         )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,5 @@
 torch==2.8.0
 torchvision==0.23.0
 torchaudio==2.8.0
+diffusers
 -r https://raw.githubusercontent.com/comfyanonymous/ComfyUI/master/requirements.txt