Native safetensors recommender and SD 2.1 hosting pipeline#29
Merged
Conversation
recommend(model, vram_gb, st_caps) returns {precision, devices, offload,
max_pixels, fork_suggested, note} with per-model tier ladders. nf4 is the
default tier; fp8/bf16 upgrade in when the card fits AND the installed
safetensors can READ the dtype. .st_can_read is a read probe distinct
from the existing write probe (CRAN safetensors reads bf16 it cannot
write, so the write probe is the wrong signal for a hosted bf16
artifact); it hand-builds a tiny safetensors file with writeBin and
loads it back. When a higher tier fits but is unreadable, recommend()
returns nf4 and surfaces the fork suggestion in $note instead of erroring.
flux_memory_profile now delegates to recommend("flux1"), fixing the stale
tiers that put fp8 (GPU-resident now) in a narrow low-VRAM band it can no
longer fit. flux_load_pipeline reconciles the recommendation with the
on-disk artifact via .flux_resolve_precision, so a widened fp8
recommendation never points at an unbuilt artifact.
R safetensors overflows a 32-bit offset on files at/above 2^31 bytes, so the old 4e9 shard default produced shards stock CRAN safetensors cannot read back (only the cornball-ai fork could). flux_quantize and the LTX nf4/fp8 writers now default shard_bytes to 1.9e9, keeping every shard under the ~2.15 GB ceiling with headroom for the post-add flush overshoot. 4e9 stays available via the argument for local fork builds.
When a user explicitly requests fp8/bf16 and the safetensors capability is missing, .st_graceful_precision prints the fork suggestion and falls back to nf4 instead of erroring; wired into download_flux1/flux2/zimage. The low-level flux_quantize/flux_load_transformer keep their hard errors (builder primitives, and a test pins that). Checkpoint reads route through .st_read_or_breadcrumb: a read that fails on a shard at/above 2^31 bytes is retranslated from the cryptic 32-bit overflow into an actionable rebuild-with-smaller-shards-or-install-fork message. Reactive (only on a real failure of an oversize shard), so it never false-alarms on the fork or on sub-2 GB shards.
…cally Documents recommend(), the small-shard rationale, the graceful fallback, and the read breadcrumb. Corrects the impression that #28 alone made nf4 load on stock CRAN safetensors: the sub-2 GB shard default is what does it, since R safetensors overflows on files at/above 2^31 bytes.
load_unet_safetensors / load_unet_sdxl_safetensors map the diffusers
UNet2DConditionModel state-dict keys 1:1 into the native modules, the
only rename being time_embedding.linear_{1,2} -> time_embedding_linear_
{1,2} (SDXL adds add_embedding.linear_{1,2}). Reuses .flux_open_sharded_dir
so single-file and sharded checkpoints both work and the >2 GB read
breadcrumb comes along; verifies every native parameter is filled and no
key or shape is left unmatched.
Validated against the real cached SDXL base UNet: all 1680 keys map with
matching shapes. The VAE decoder and CLIP text encoder already had
safetensors loaders; the UNet was the gap. Portable test drives the
copy/completeness/error paths through a mock module.
unet_native_from_safetensors / unet_sdxl_native_from_safetensors build the native UNet and load its weights from a diffusers directory, the safetensors counterparts to the *_from_torchscript constructors (no TorchScript, so Blackwell-safe). Thin wrappers over the validated default construction plus the safetensors loader; ... forwards variant overrides and the loader's shape check fails loudly on a mismatched architecture. Validated end-to-end against the cached SDXL base UNet: 1680 params built + loaded in ~20 s. Opt-in gated test documents it.
vae_decoder_native_from_safetensors and text_encoder_native_from_safetensors build and load their native modules from a diffusers directory (no TorchScript, Blackwell-safe). The text-encoder builder reads CLIPTextConfig from config.json via .detect_text_encoder_config (the config counterpart to the TorchScript arch detector). Validated against cached FLUX analogs: the FLUX VAE exercises the shared 16-channel decoder, and FLUX's text_encoder IS the SDXL CLIP ViT-L (config-detected 768/12/12, forward -> 1x77x768).
download_sd21() fetches the SD 2.1 diffusers weights (gated-repo breadcrumb if the license is unaccepted); sd_pipeline_from_safetensors() assembles the native UNet + VAE decode + CLIP text encoder from a diffusers directory, and txt2img_sd21(diffusers_dir=) runs it with no TorchScript (Blackwell-safe). Fixes the SD VAE decode: the native decoder (derived from the FLUX VAE, which has no quant convs) omitted post_quant_conv, so decode(z) skipped the 1x1 conv diffusers applies before the decoder -- cos ~0.5 vs the reference, i.e. garbage. sd_vae_decode now applies it from the artifact's own post_quant_conv weights (cos 1.0). SD 2.1 defaults to float32 here (fp16 attention overflows to NaN). Known issue (pre-existing, tracked): the native SD 2.1 UNet has a spatial-coherence "tiling" bug that the .pt-native path shares (identical weights); the existing parity test misses it by feeding constant inputs.
The SD 2.1 native UNet computed its timestep embedding with flip_sin_to_cos=FALSE, downscale_freq_shift=1, which scrambles the sin/cos channel ordering the trained time_embedding weights expect. Standard diffusers SD (and the native SDXL UNet) use TRUE/0. The error was invisible to the constant-input parity test (GroupNorm sits at its bias and attention is uniform, so the spatial path is never exercised) but compounded through the network into spatially-tiled output - recognizable content repeated across the frame instead of one subject. Localized by comparing every native submodule to its callable TorchScript counterpart (all cos 1.0), then reconstructing the forward from TS submodules and sweeping wiring params: flip=TRUE,shift=0 -> cos 1.000000 vs the reference (was 0.82 at 64x64). Native UNet now matches TorchScript at cos 0.99999 across resolutions; SD 2.1 renders a coherent subject instead of tiles. test_unet.R gains a random-input parity check so this class of bug can't hide behind constant inputs again.
convert_sd21_pt_to_diffusers() rebuilds a diffusers-layout safetensors directory (unet/, vae/, text_encoder/ + derived CLIPTextConfig) from the cornball-ai/sd21-R TorchScript .pt files. A trace preserves the exact parameter tensors, so the output is bit-identical to the source at the chosen dtype (verified max|Δ|=0 vs the .pt). float16 keeps every file sub-2 GB (CRAN-safetensors readable). This is the provenance-clean build step for the hosted artifact now that upstream SD 2.1 was deprecated; cornball already hosts these OpenRAIL weights as .pt.
The upstream stabilityai/stable-diffusion-2-1 repo was deprecated; download_sd21 now fetches the fp16 diffusers safetensors from the cornball-ai/sd21-R dataset (diffusers/ prefix, repo_type="dataset"), built from our own .pt via convert_sd21_pt_to_diffusers and uploaded with hub_upload. No unet/vae config.json is needed (the native constructors use the SD 2.1 defaults); only the CLIP encoder config ships.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Phase 0 foundations plus the SD 2.1 safetensors pipeline from the uniform-hosting plan.
Recommender and capability probing
recommend(model, vram, st_caps): one VRAM and safetensors-read-capability policy for every model. nf4 is the default tier; fp8/bf16 come in when the card fits them and the installed safetensors can read the dtype, otherwise it recommends nf4 and surfaces the fork suggestion..st_can_read()read probe, distinct from the write probe (CRAN reads bf16 it cannot write).flux_memory_profile()now delegates torecommend(), fixing the stale tiers.Shards and graceful degradation
SD native from safetensors
*_from_safetensorsconstructors for the SD/SDXL UNet, VAE decoder, and CLIP text encoder (no TorchScript, Blackwell-safe). SDXL UNet validated against the cached base weights, 1680/1680 keys.download_sd21()+sd_pipeline_from_safetensors()run SD 2.1 end to end from diffusers safetensors, pointed at the hostedcornball-ai/sd21-Rartifact.convert_sd21_pt_to_diffusers()builds that artifact from our TorchScript.pt(bit-identical), the provenance-clean source now that upstream SD 2.1 was deprecated.Two fixes the pipeline surfaced
post_quant_conv(cos ~0.5 vs the reference); restored to cos 1.0.Version 0.1.0.6. Full suite green (805 tests).