import importlib.util import os import random import shutil import subprocess import sys import gradio as gr import spaces import torch from huggingface_hub import hf_hub_download, snapshot_download from PIL import Image # -------------------------------------------------------------------------------------- # 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. # -------------------------------------------------------------------------------------- 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" ) DEFAULT_NEGATIVE = ( "oldest, low quality, lowres, blurry, out of focus, jpeg artifacts, watermark, signature, " "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, ) 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"), ) SUPRA_CODE = hf_hub_download(SUPRA_REPO, "inference.py") SUPRA_CKPT = hf_hub_download(SUPRA_REPO, "model_final_ema.pt") _rt = {} # ------------------------------------- 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, 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, 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 _rt.update(nodes=nodes, Loader=Nanosaur2Loader, LoaderInt8=Nanosaur2LoaderInt8) return _rt def snap(v, m=16): return max(m, int(round(v / m)) * m) def run_nanosaur(int8, prompt, negative, width, height, steps, cfg, guidance_mode, seed): rt = get_comfy() nodes = rt["nodes"] 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 = loader().load( unet, "nanosaur2_text_encoder.safetensors", "nanosaur2_vae.safetensors", guidance_mode ) 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", 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") 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( "# 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(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 (Nanosaur2 only)", lines=2, value=DEFAULT_NEGATIVE) run = gr.Button("Generate", variant="primary", size="lg") 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") 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") 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(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, [models, prompt, negative_prompt, width, height, steps, cfg, guidance_mode, supra_steps, supra_cfg, seed, randomize_seed], [gallery, used_seed], ) if __name__ == "__main__": demo.launch(theme=gr.themes.Soft(), css=CSS)