mirror of
https://github.com/vladmandic/automatic
synced 2026-08-28 08:00:59 +02:00
Dity0
101 lines
4.4 KiB
Python
101 lines
4.4 KiB
Python
#!/usr/bin/env python
|
|
"""Standalone end-to-end smoke for diffusers-native Ideogram 4.
|
|
|
|
Loads the split-projection bf16 diffusers folder and quantizes both transformers
|
|
with SDNQ at load (the same path SD.Next uses), loads the shared Qwen3-VL text
|
|
encoder, builds the diffusers ``Ideogram4Pipeline``, and generates an image. This
|
|
exercises the real pipeline: both transformers under SDNQ, the Qwen3-VL 13-layer
|
|
tap, the dual-branch asymmetric CFG loop, the logit-normal schedule, the vae.bn
|
|
latent denorm, and VAE decode. SDNQ int4 fits the two towers plus the encoder on a
|
|
24GB GPU.
|
|
|
|
Usage:
|
|
python test/test-ideogram4-smoke.py --model /path/to/Ideogram-4-bf16-split --output out.png
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", required=True, help="split-projection bf16 diffusers folder")
|
|
parser.add_argument("--output", default="ideogram4_smoke.png")
|
|
parser.add_argument("--prompt", default="a ginger cat wearing a tiny wizard hat reading a glowing spellbook, detailed digital illustration")
|
|
parser.add_argument("--height", type=int, default=1024)
|
|
parser.add_argument("--width", type=int, default=1024)
|
|
parser.add_argument("--steps", type=int, default=20)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--weights-dtype", default="uint4", help="SDNQ weights dtype (uint4, int8, ...)")
|
|
parser.add_argument("--hf-cache", default=None, help="HF cache_dir for the shared Qwen3-VL encoder (default: HF default cache)")
|
|
args = parser.parse_args()
|
|
|
|
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, REPO_ROOT)
|
|
os.chdir(REPO_ROOT)
|
|
os.environ["SD_INSTALL_QUIET"] = "1"
|
|
|
|
# Our own args are already parsed; clear argv (and leave it cleared) so sdnext's
|
|
# shared.py / devices, which re-parse argv on import, don't see this script's flags.
|
|
sys.argv = [sys.argv[0]]
|
|
|
|
import modules.cmd_args
|
|
import installer
|
|
|
|
modules.cmd_args.parse_args()
|
|
installer.add_args(modules.cmd_args.parser)
|
|
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
|
|
|
|
import torch
|
|
import diffusers
|
|
from transformers import AutoTokenizer
|
|
from transformers.models.qwen3_vl import Qwen3VLModel
|
|
|
|
from modules import devices
|
|
from sdnq import SDNQConfig
|
|
|
|
TE_REPO = "Qwen/Qwen3-VL-8B-Instruct"
|
|
|
|
|
|
def main() -> int:
|
|
device = devices.device
|
|
cfg = SDNQConfig(weights_dtype=args.weights_dtype)
|
|
|
|
print(f"loading transformer (sdnq {args.weights_dtype}) ...", flush=True)
|
|
transformer = diffusers.Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
|
|
print("loading unconditional_transformer ...", flush=True)
|
|
uncond = diffusers.Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="unconditional_transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
|
|
print("loading text encoder Qwen3-VL ...", flush=True)
|
|
te_kwargs = {"cache_dir": args.hf_cache} if args.hf_cache else {}
|
|
text_encoder = Qwen3VLModel.from_pretrained(TE_REPO, quantization_config=SDNQConfig(weights_dtype=args.weights_dtype), torch_dtype=torch.bfloat16, **te_kwargs).to(device)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model, subfolder="tokenizer")
|
|
print("loading vae ...", flush=True)
|
|
vae = diffusers.AutoencoderKLFlux2.from_pretrained(args.model, subfolder="vae", torch_dtype=torch.bfloat16).to(device)
|
|
scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(args.model, subfolder="scheduler")
|
|
|
|
pipe = diffusers.Ideogram4Pipeline(
|
|
scheduler=scheduler,
|
|
vae=vae,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
transformer=transformer,
|
|
unconditional_transformer=uncond,
|
|
)
|
|
|
|
generator = torch.Generator(device=device).manual_seed(args.seed)
|
|
print(f"generating {args.width}x{args.height} steps={args.steps} ...", flush=True)
|
|
start = time.time()
|
|
out = pipe(prompt=args.prompt, num_inference_steps=args.steps, guidance_scale=7.0, guidance_schedule=None, width=args.width, height=args.height, generator=generator)
|
|
elapsed = time.time() - start
|
|
|
|
image = out.images[0]
|
|
image.save(args.output)
|
|
if torch.cuda.is_available():
|
|
print(f"peak VRAM: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB", flush=True)
|
|
print(f"PASS: generated {args.output} in {elapsed:.1f}s, size={image.size}", flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|