Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-02 11:00:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4994 from vladmandic/dev

Browse Source 
merge dev
This commit is contained in:
Vladimir Mandic
2026-07-14 10:24:45 +02:00
committed by GitHub
parent 2763185abd 7d9130aef0
commit e08063a963
111 changed files with 8410 additions and 1119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/skills/check-processing/SKILL.md
+5 -3
View File
@@ -27,7 +27,7 @@ Run workflow coverage in this order to keep checks focused and complete:
2. `img2img`: `modules/ui_img2img.py` -> `modules/img2img.py` -> `modules/processing.py:process_images` -> `modules/processing_diffusers.py:process_diffusers`
3. `control/process`: `modules/ui_control.py` -> `modules/control/run.py` (and related control processing entrypoints) -> `modules/processing.py:process_images` -> `modules/processing_diffusers.py:process_diffusers`
4. `caption/process`: `modules/ui_caption.py` -> caption handler module(s) -> `modules/processing.py:process_images` and/or postprocess/caption execution module(s), depending on selected caption backend
5. `video`: `modules/ui_video.py` -> `modules/video_models/video_run -> `modules/processing.py:process_images` and/or postprocess/video execution module(s), depending on implementation
5. `video`: `modules/ui_video.py` -> `modules/video_models/video_run.py` -> `modules/processing.py:process_images` and/or postprocess/video execution module(s), depending on implementation
Also validate script hooks when present:
@@ -40,8 +40,10 @@ Also validate script hooks when present:
- `modules/ui_img2img.py`
- `modules/img2img.py`
- `modules/ui_control.py`
- `modules/ui_caption.py` (and `modules/ui_captions.py` if present)
- `modules/control/run.py`
- `modules/ui_caption.py`
- `modules/ui_video.py`
- `modules/video_models/video_run.py`
- `modules/processing.py`
- `modules/processing_diffusers.py`
- `modules/scripts_manager.py`
@@ -58,7 +60,7 @@ For each covered workflow, verify all three dimensions:
### 1. Build End-To-End Call Graph
For each workflow (`txt2img`, `img2img`, `control`, `caption`):
For each workflow (`txt2img`, `img2img`, `control`, `caption`, `video`):
- Locate submit/click bindings in UI modules.
- Capture the exact `inputs=[...]` list order and target function (`fn=...`).
.gitignore
+3 -1
View File
@@ -6,7 +6,6 @@ __pycache__
.idea/
**/.DS_Store
/*.json
/*.yaml
/params.txt
/styles.csv
/user.css
@@ -18,6 +17,8 @@ __pycache__
/data/themes.json
/data/installer.json
/data/rocm.json
/data/signatures.json
/data/civitai.json
node_modules
package-lock.json
@@ -60,6 +61,7 @@ tunableop_results*.csv
!requirements.txt
!constraints.txt
!pnpm-lock.yaml
!pnpm-workspace.yaml
!/data
!/models/VAE-approx
!/models/VAE-approx/model.pt
.vscode/settings.json
Vendored
+2 -1
View File
@@ -20,5 +20,6 @@
},
"githubPullRequests.ignoredPullRequestBranches": [
"master"
]
],
"chat.byokUtilityModelDefault": "mainAgent"
}
CHANGELOG.md
+49 -2
View File
@@ -1,9 +1,49 @@
# Change Log for SD.Next
## Update for 2026-07-08
## Update for 2026-07-14
### Highlights for 2026-07-14
*What's New?* Full week(!) since there release, we're bringing a service pack update:
**Anima** has new *Aesthetic* and *Turbo* variants, **Joy Image Edit** has new *Plus* variant
*And also*:
- UI updates to *Server info* and *Log viewer*, more informative and allows easier sharing of info
- Quite a few *SDNQ* improvements, both to quantization and attention
- Continuing with enhanced support for 3rd party finetunes
*Note*: As long as there are no major items in `#dev` branch, we will continue with weekly service updates, so stay tuned for more!
[Home](https://vladmandic.github.io/sdnext/) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic)
### Details for 2026-07-14
- **Models**
- [Anima 1.0](https://civitai.com/models/2458426/anima?modelVersionId=3108569) new *Aesthetic* and *Turbo* variants
- [Joy Image Edit](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus-Diffusers) new *Plus* variant that has multi-image-edit capabilities
- **Features**
- api add `/sdapi/v1/restart` endpoint
- loader: support loading secondary unet for models that have multiple unets
*networks panel -> toggle: load model as secondary model*
*settings -> model loading -> unet model secondary*
- loader: support transformers from all-in-one safetensors
- loader: improve `fp8` model support
- loader: improve `nvfp4` model support
- lora: additional lora types
- civitai: improve search
- civitai: restart download
- sdnq: improve attention
- api: add `/sdapi/v1/restart` endpoint
- server: strict sub-path mount and redirects for all endpoints
- **UI**
- improved server info panel
now shows loaded components, networks, etc.
- improved log viewer for both server and client logs
- expose server errors
- log view copy server log and copy client log buttons
copies current log to clipboard for easy sharing
- collapsible input & output panels for *video*, *caption*, *process* tabs
- video add immediate load
- **Experimental**
- support for `comfy_quant` models
- **Fixes**
- huggingface: fix model access on windows
- lora: fix lora text-encoder loader
@@ -11,6 +51,13 @@
- api: validate swagger schema
- installer: handle detached branch on non-english setups
- fs: faster stat
- server-info: handle multiple gpus
- load: skip redundant model reload
- sdnq: fix hadamard on sdnq atten with sd15
- download: improved progress tracking
- gallery: delete image
- gallery: save image
- standardui: gallery collapse sections
## Update for 2026-07-07
TODO.md
+4 -2
View File
@@ -8,7 +8,6 @@
### Assigned
- [Object clear](https://huggingface.co/jixin0101/ObjectClear) remover for Kanvas, @vladmandic
- Chat-based interface, @vladmandic
- Control tab verify overrides handling, @vladmandic
- Cloud providers, @CalamitousFelicitousness
@@ -17,10 +16,13 @@
### Unassigned
- [Nunchaku Lite](https://github.com/huggingface/diffusers/pull/14100)
- Processing -> Video capabilities
- `RIFE` in processing
- `SeedVR2` in processing
- Video model loader: Add video models to Reference
- [Object clear](https://huggingface.co/jixin0101/ObjectClear) remover for Kanvas
- Video models: add to Reference
- Video models: support custom entries
- UI Lite vs Expert mode
- Auto handle scheduler `prediction_type`
- Cache models in memory
cli/sdnq-attention-benchmark.py
Executable
+3025
View File
File diff suppressed because it is too large Load Diff
data/previews.json
+29 -29
View File
@@ -1,41 +1,41 @@
{
"sd-v21-512-ema": "stabilityai--stable-diffusion-2-1-base.jpg",
"stable-diffusion-v1-5--stable-diffusion-v1-5": "v1-5-pruned-fp16-emaonly.jpg",
"stabilityai--stable-diffusion-xl-base-1.0": "stabilityai--stable-diffusion-xl-base-1.0.jpg",
"stabilityai--stable-diffusion-3-medium-diffusers": "stabilityai--stable-diffusion-3.jpg",
"stabilityai--stable-diffusion-3.5-medium": "stabilityai--stable-diffusion-3_5-medium.jpg",
"stabilityai--stable-diffusion-3.5-large": "stabilityai--stable-diffusion-3_5-large.jpg",
"stabilityai--stable-diffusion-3.5-large-turbo": "stabilityai--stable-diffusion-3_5-large-turbo.jpg",
"Disty0--FLUX.1-dev-qint8": "black-forest-labs--FLUX.1-dev.jpg",
"chroma-unlocked-v50-annealed": "lodestones--Chroma-annealed.jpg",
"chroma-unlocked-v50": "lodestones--Chroma-detail.jpg",
"Disty0--FLUX.1-dev-qint4": "black-forest-labs--FLUX.1-dev.jpg",
"sayakpaul--flux.1-dev-nf4": "black-forest-labs--FLUX.1-dev.jpg",
"THUDM--CogVideoX-2b": "THUDM--CogView3-Plus-3B.jpg",
"THUDM--CogVideoX-5b": "THUDM--CogView3-Plus-3B.jpg",
"THUDM--CogVideoX-5b-I2V": "THUDM--CogView3-Plus-3B.jpg",
"Disty0--FLUX.1-dev-qint8": "black-forest-labs--FLUX.1-dev.jpg",
"Efficient-Large-Model--Sana_1600M_1024px_BF16_diffusers": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"Efficient-Large-Model--Sana_1600M_2Kpx_BF16_diffusers": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"Efficient-Large-Model--Sana_1600M_4Kpx_BF16_diffusers": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"Efficient-Large-Model--Sana_600M_1024px_diffusers": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"stabilityai--stable-video-diffusion-img2vid-xt-1-1": "stabilityai--stable-video-diffusion-img2vid-xt.jpg",
"shuttleai--shuttle-3-diffusion": "shuttleai--shuttle-3-diffusion.jpg",
"HiDream-I1-Full": "HiDream-I1 Full",
"lodestones--Chroma1-Base": "lodestones--Chroma-Base.jpg",
"lodestones--Chroma1-HD": "lodestones--Chroma-HD.jpg",
"chroma-unlocked-v50": "lodestones--Chroma-detail.jpg",
"chroma-unlocked-v50-annealed": "lodestones--Chroma-annealed.jpg",
"vladmandic--Qwen-Lightning": "Qwen-Lightning.jpg",
"vladmandic--Qwen-Lightning-Edit": "Qwen-Lightning.jpg",
"Wan-AI--Wan2.2-T2V-A14B-Diffusers": "Wan-AI--Wan2.2-T2V-A14B-Diffusers.jpg",
"Wan-AI--Wan2.1-T2V-14B-Diffusers": "Wan-AI--Wan2.1-T2V-14B-Diffusers.jpg",
"linoyts--Wan2.2-VACE-Fun-14B-diffusers": "linoyts--Wan2.2-VACE-Fun-14B-diffusers.jpg",
"vladmandic--Anima-1.0-Base-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Base.jpg",
"vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Turbo.jpg",
"vladmandic--Krea-2-Turbo-sdnq-hadamard-uint4": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
"vladmandic--Krea-2-Base-sdnq-hadamard-uint4": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
"Jinstudio--Lens": "microsoft--Lens.jpg",
"Jinstudio--Lens-Base": "microsoft--Lens-Base.jpg",
"Jinstudio--Lens-Turbo": "microsoft--Lens-Turbo.jpg",
"Jinstudio--Lens": "microsoft--Lens.jpg",
"linoyts--Wan2.2-VACE-Fun-14B-diffusers": "linoyts--Wan2.2-VACE-Fun-14B-diffusers.jpg",
"lodestones--Chroma1-Base": "lodestones--Chroma-Base.jpg",
"lodestones--Chroma1-HD": "lodestones--Chroma-HD.jpg",
"SahilCarterr--BRIA-3.2": "briaai--BRIA-3.2.jpg",
"sayakpaul--flux.1-dev-nf4": "black-forest-labs--FLUX.1-dev.jpg",
"sd-v21-512-ema": "stabilityai--stable-diffusion-2-1-base.jpg",
"shuttleai--shuttle-3-diffusion": "shuttleai--shuttle-3-diffusion.jpg",
"stabilityai--stable-diffusion-3-medium-diffusers": "stabilityai--stable-diffusion-3.jpg",
"stabilityai--stable-diffusion-3.5-large-turbo": "stabilityai--stable-diffusion-3_5-large-turbo.jpg",
"stabilityai--stable-diffusion-3.5-large": "stabilityai--stable-diffusion-3_5-large.jpg",
"stabilityai--stable-diffusion-3.5-medium": "stabilityai--stable-diffusion-3_5-medium.jpg",
"stabilityai--stable-diffusion-xl-base-1.0": "stabilityai--stable-diffusion-xl-base-1.0.jpg",
"stabilityai--stable-video-diffusion-img2vid-xt-1-1": "stabilityai--stable-video-diffusion-img2vid-xt.jpg",
"stable-diffusion-v1-5--stable-diffusion-v1-5": "v1-5-pruned-fp16-emaonly.jpg",
"THUDM--CogVideoX-2b": "THUDM--CogView3-Plus-3B.jpg",
"THUDM--CogVideoX-5b-I2V": "THUDM--CogView3-Plus-3B.jpg",
"THUDM--CogVideoX-5b": "THUDM--CogView3-Plus-3B.jpg",
"vladmandic--Anima-1.0-Base-Merge-sdnq-hadamard-uint4": "vladmandic--Anima-1.0-Base.jpg",
"vladmandic--Anima-1.0-Base-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Base.jpg",
"vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Turbo.jpg",
"vladmandic--Flux.2-Klein-9B-KV-sdnq-hadamard-uint4": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
"vladmandic--Anima-1.0-Base-Merge-sdnq-hadamard-uint4": "vladmandic--Anima-1.0-Base.jpg"
"vladmandic--Krea-2-Base-sdnq-hadamard-uint4": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
"vladmandic--Krea-2-Turbo-sdnq-hadamard-uint4": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
"vladmandic--Qwen-Lightning-Edit": "Qwen-Lightning.jpg",
"vladmandic--Qwen-Lightning": "Qwen-Lightning.jpg",
"Wan-AI--Wan2.1-T2V-14B-Diffusers": "Wan-AI--Wan2.1-T2V-14B-Diffusers.jpg",
"Wan-AI--Wan2.2-T2V-A14B-Diffusers": "Wan-AI--Wan2.2-T2V-A14B-Diffusers.jpg"
}
data/reference-base.json
+16
View File
@@ -288,6 +288,14 @@
"date": "2026 May",
"size": 4.99
},
"Anima 1.0 Aesthetic": {
"path": "CalamitousFelicitousness/Anima-1.0-Aesthetic-Diffusers",
"preview": "CalamitousFelicitousness--Anima-1.0-Aesthetic-Diffusers.jpg",
"desc": "Anima 1.0 Aesthetic, fine-tuned for better consistency and a higher quality default art style. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
"extras": "sampler: Default, cfg_scale: 4.5, steps: 30",
"date": "2026 July",
"size": 4.99
},
"Meituan LongCat Image": {
"path": "meituan-longcat/LongCat-Image",
"preview": "meituan-longcat--LongCat-Image.jpg",
@@ -902,5 +910,13 @@
"size": 46.86,
"extras": "sampler: Default",
"date": "2026 April"
},
"JoyAI Image Edit Plus": {
"path": "jdopensource/JoyAI-Image-Edit-Plus-Diffusers",
"preview": "jdopensource--JoyAI-Image-Edit-Diffusers.jpg",
"desc": "JoyAI Image Edit Plus is a Diffusers-native image editing model that combines a JoyImageEditPlus transformer with Qwen3-VL multimodal conditioning for instruction-guided edits.",
"size": 47.98,
"extras": "sampler: Default",
"date": "2026 July"
}
}
data/reference-distilled.json
+5 -4
View File
@@ -212,10 +212,11 @@
"date": "2026 March"
},
"Anima 1.0 Turbo": {
"path": "vladmandic/Anima-1.0-Turbo",
"preview": "vladmandic--Anima-1.0-Turbo.jpg",
"desc": "Anima 1.0 Turbo with extended 1024-resolution training and expanded dataset coverage for less common artists. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
"date": "2026 May",
"path": "CalamitousFelicitousness/Anima-1.0-Turbo-Diffusers",
"preview": "CalamitousFelicitousness--Anima-1.0-Turbo-Diffusers.jpg",
"desc": "Anima 1.0 Turbo, distilled for fast generation with increased stability and a strong default style. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 10",
"date": "2026 July",
"size": 4.99
},
"Meituan LongCat Image-Edit Turbo": {
extensions-builtin/sdnext-kanvas
+1 -1
Submodule extensions-builtin/sdnext-kanvas updated: 4f650445b0...0261cc2830
extensions-builtin/sdnext-modernui
+1 -1
Submodule extensions-builtin/sdnext-modernui updated: 4b8d5f6778...07890f0d7b
installer.py
+1 -1
View File
@@ -549,7 +549,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all:
return
target_commit = "ea802951f5fb235b6af8fe9247f56187d49748b2" # diffusers commit hash == 0.39.0.dev0 == 06-29-2026
target_commit = "01969142b55379991fee07608c9e7e8f80afced0" # diffusers commit hash == 0.39.0.dev0 == 06-29-2026
# if args.use_rocm or args.use_zluda or args.use_directml:
# sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
pkg = package_spec('diffusers')
models/Reference/CalamitousFelicitousness--Anima-1.0-Aesthetic-Diffusers.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 60 KiB

models/Reference/CalamitousFelicitousness--Anima-1.0-Turbo-Diffusers.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 96 KiB

models/Reference/vladmandic--Anima-1.0-Turbo.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 42 KiB

modules/api/api.py
+1 -1
View File
@@ -58,7 +58,7 @@ class Api:
self.add_api_route("/sdapi/v1/restart", server.post_restart, methods=["POST"], status_code=204, tags=["Server"])
self.add_api_route("/sdapi/v1/memory", server.get_memory, methods=["GET"], response_model=models.ResMemory, tags=["Server"])
self.add_api_route("/sdapi/v1/cmd-flags", server.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel, tags=["Server"])
self.add_api_route("/sdapi/v1/gpu", gpu.get_gpu, methods=["GET"], tags=["Server"])
self.add_api_route("/sdapi/v1/gpu", gpu.get_gpu, methods=["GET"], tags=["Server"], response_model=list[dict])
self.add_api_route("/sdapi/v1/gpu-smi", gpu.get_gpu_smi, methods=["GET"], response_model=list[models.ResGPU], tags=["Server"])
# core api using locking
modules/api/gpu.py
+1 -6
View File
@@ -7,12 +7,7 @@ device = None
def get_gpu():
import installer
res = {}
if len(installer.gpu_info) == 1:
return installer.gpu_info[0]
for i, item in enumerate(installer.gpu_info):
res[i] = item
return res
return installer.gpu_info
def get_gpu_smi():
modules/api/loras.py
+7
View File
@@ -26,6 +26,12 @@ def post_refresh_loras():
return result
def get_loaded():
"""Return a list of currently loaded LoRA networks with name, alias, path, and metadata."""
from modules.lora import lora_common
return [obj.name for obj in lora_common.loaded_networks]
def _invalidate_extra_networks():
"""Reset extra-networks page caches so the v2 API picks up changes."""
from modules import shared
@@ -38,3 +44,4 @@ def register_api(api):
api.add_api_route("/sdapi/v1/lora", get_lora, methods=["GET"], response_model=dict, tags=["Enumerators"])
api.add_api_route("/sdapi/v1/loras", get_loras, methods=["GET"], response_model=list[dict], tags=["Enumerators"])
api.add_api_route("/sdapi/v1/refresh-loras", post_refresh_loras, methods=["POST"], tags=["Functional"])
api.add_api_route("/sdapi/v1/loaded-loras", get_loaded, methods=["GET"], response_model=list[str], tags=["Enumerators"])
modules/api/middleware.py
+19
View File
@@ -9,6 +9,7 @@ import fastapi
from starlette.responses import JSONResponse
from fastapi import FastAPI, Request, Response
from fastapi.exceptions import HTTPException
from fastapi.responses import RedirectResponse
from fastapi.encoders import jsonable_encoder
from modules.logger import log
import modules.errors as errors
@@ -17,6 +18,19 @@ from modules.api.validate import validate_request, validate_log
errors.install()
def validate_subpath(endpoint: str, subpath: str | None):
if (subpath is not None) and (len(subpath) > 0) and (not endpoint.startswith(subpath)):
if not subpath.endswith('/'):
subpath = f'{subpath}/'
if endpoint.startswith('/'):
url = f'{subpath}{endpoint[1:]}'
else:
url = f'{subpath}{endpoint}'
log.trace(f'API: redirect subpath={subpath} url="{endpoint}" redirect="{url}"')
return RedirectResponse(url=url, status_code=308)
return None
def setup_middleware(app: FastAPI, cmd_opts):
ssl._create_default_https_context = ssl._create_unverified_context # pylint: disable=protected-access
uvicorn_logger=logging.getLogger("uvicorn.error")
@@ -43,6 +57,11 @@ def setup_middleware(app: FastAPI, cmd_opts):
endpoint = req.scope.get('path', 'err')
client = req.scope.get('client', ('0:0.0.0', 0))[0]
token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure")
redirect = validate_subpath(endpoint, cmd_opts.subpath)
if redirect:
return redirect
validate_request(client, endpoint)
if cmd_opts.api_log:
if not validate_log(client, endpoint):
modules/api/validate.py
+14 -10
View File
@@ -31,6 +31,7 @@ log_cost = {
"/sdapi/v1/memory": -1,
"/sdapi/v1/platform": -1,
"/sdapi/v1/checkpoint": -1,
"/sdapi/v1/loaded-loras": -1,
"/sdapi/v1/gpu-smi": -1,
"/sdapi/v1/status": 60,
"/sdapi/v1/progress": 60,
@@ -38,8 +39,9 @@ log_cost = {
log_exclude_suffix = ['.css', '.js', '.ico', '.svg']
log_exclude_prefix = ['/assets']
class Limiter():
def __init__(self, limit):
def __init__(self, limit, subpath=None):
import limits
self.request_backend = limits.storage.MemoryStorage()
self.request_limit = limit # default is 300 requests per minute
@@ -50,7 +52,8 @@ class Limiter():
self.log_strategy = limits.strategies.FixedWindowRateLimiter(self.log_backend)
self.log_limiter = limits.parse(f"{self.log_limit}/minute")
self.summary = {}
log.info(f'API: limit={self.request_limit} strategy={self.request_strategy.__class__.__name__} backend={self.request_backend.__class__.__name__}')
self.subpath = subpath
log.info(f'API: limit={self.request_limit} strategy={self.request_strategy.__class__.__name__} backend={self.request_backend.__class__.__name__} subpath={self.subpath}')
def stats(self):
@@ -66,8 +69,8 @@ class Limiter():
return False
status = self.request_strategy.hit(self.request_limiter, client, api, cost=cost)
if not status and not quiet:
log.warning(f'API: client={client} api={api} rate limit exceeded')
from fastapi.exceptions import HTTPException
log.warning(f'API: client={client} api={api} rate limit exceeded')
raise HTTPException(status_code=429, detail=f"{client}:{api}: rate limit exceeded")
return status
@@ -85,7 +88,7 @@ class Limiter():
return status
limiter = Limiter(300)
limiter = Limiter(0, None)
def get_api_stats():
@@ -94,20 +97,21 @@ def get_api_stats():
def validate_request(client, endpoint):
global limiter # pylint: disable=global-statement
from modules.shared import opts
from modules.shared import opts, cmd_opts
if opts.server_rate_limit != limiter.request_limit:
limiter = Limiter(opts.server_rate_limit)
limiter = Limiter(opts.server_rate_limit, cmd_opts.subpath)
api = re.match(r"^[^?#&=]+", endpoint).group(0)
if (limiter.subpath is not None) and (len(limiter.subpath) > 0) and api.startswith(limiter.subpath): # strip subpath from api for rate limiting
api = api[len(limiter.subpath):]
key = f"{client}:{api}"
if key not in limiter.summary:
limiter.summary[key] = 0
limiter.summary[key] += 1
# import anyio
# _limiter = anyio.to_thread.current_default_thread_limiter()
# log.debug(f'FastAPI: threads={_limiter._total_tokens}')
return limiter.check_request(client, api)
def validate_log(client, endpoint):
api = re.match(r"^[^?#&=]+", endpoint).group(0)
if (limiter.subpath is not None) and (len(limiter.subpath) > 0) and api.startswith(limiter.subpath): # strip subpath from api for logging
api = api[len(limiter.subpath):]
return limiter.check_log(client, api)
modules/attention.py
+5 -1
View File
@@ -24,7 +24,11 @@ def set_sdnq_attention():
sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sdnq_atten)
def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
if query.device.type != "cpu" and query.shape[-3] > 1: # Skip VAE
if (
query.device.type != "cpu"
and (query.shape[-2] >= 512 or key.shape[-2] >= 512) # Skip TE
and query.shape[-3] > 1 # Skip VAE
):
return sdnq_triton_atten(
query=query, key=key, value=value, attn_mask=attn_mask,
is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
modules/civitai/api_civitai.py
+113 -2
View File
@@ -103,10 +103,31 @@ def get_search(
base_models=bm_list, nsfw=nsfw, limit=limit, cursor=cursor,
username=username, favorites=favorites, token=token,
)
history_params = {k: v for k, v in {
'types': types,
'sort': sort,
'period': period,
'base_models': base_models,
'nsfw': nsfw,
'username': username,
'favorites': favorites,
}.items() if v}
if query:
search_history.add('query', query)
search_history.add('query', query, history_params)
elif tag:
search_history.add('tag', tag)
search_history.add('tag', tag, history_params)
elif history_params:
# Filter-only browse: label the entry with the filters themselves.
parts = [
types,
base_models,
'favorites' if favorites else (f'by {username}' if username else ''),
period,
sort,
'nsfw' if nsfw else '',
]
label = ' · '.join(p for p in parts if p)
search_history.add('filter', label, history_params)
return response.dict(by_alias=True)
@@ -234,6 +255,95 @@ def get_download_status():
return download_manager.status()
peek_cache = None
def peek_cache_get(file_id: int, url: str):
"""Persistent probe cache keyed by civitai file id; content per id is
immutable, so entries never expire. A url hash guards against id reuse."""
global peek_cache # pylint: disable=global-statement
import hashlib
from modules import paths
from modules.json_helpers import readfile
if peek_cache is None:
peek_cache = readfile(paths.civitai_probe_file, silent=True, lock=True, as_type='dict')
entry = peek_cache.get(str(file_id))
if entry and entry.get('url_hash') == hashlib.sha256(url.encode('utf-8')).hexdigest()[:16]:
return entry.get('response')
return None
def peek_cache_put(file_id: int, url: str, response: dict):
import hashlib
from modules import paths
from modules.json_helpers import writefile
peek_cache[str(file_id)] = {
'url_hash': hashlib.sha256(url.encode('utf-8')).hexdigest()[:16],
'response': response,
}
writefile(peek_cache, paths.civitai_probe_file, silent=True, atomic=True)
def get_peek_header(url: str, file_id: int = 0):
"""Read a remote safetensors JSON header via ranged requests. Returns the
__metadata__ block plus a full model_probe analysis (architecture,
dtypes, quant scheme) without downloading the file."""
import json
import struct
from modules import shared
# civitai.red serves the same download service and rewrites downloadUrl to
# its own host; normalize so the guard, cache hash and fetch host agree.
url = url.replace('https://civitai.red/', 'https://civitai.com/', 1)
if not url.startswith('https://civitai.com/'):
return JSONResponse(content={"error": "only civitai urls are allowed"}, status_code=400)
if file_id:
cached = peek_cache_get(file_id, url)
if cached is not None:
return cached
base_headers = {}
token = getattr(shared.opts, 'civitai_token', '') or ''
if token:
base_headers['Authorization'] = f'Bearer {token}'
def read_range(start: int, end: int) -> bytes:
r = shared.req(url, headers={**base_headers, 'Range': f'bytes={start}-{end}'}, stream=True)
status = getattr(r, 'status_code', 500)
if status not in (200, 206):
raise RuntimeError(f'HTTP {status}')
want = end - start + 1
# A 200 reply means the server ignored the range and sends from byte 0.
need = end + 1 if status == 200 else want
buf = b''
try:
for chunk in r.iter_content(chunk_size=65536):
buf += chunk
if len(buf) >= need:
break
finally:
r.close()
return buf[start:start + want] if status == 200 else buf[:want]
try:
prefix = read_range(0, 7)
if len(prefix) < 8:
return {"metadata": None, "error": "short read"}
header_len = struct.unpack('<Q', prefix)[0]
if header_len <= 0 or header_len > 16 * 1024 * 1024:
return {"metadata": None, "error": f"implausible header length: {header_len}"}
header = json.loads(read_range(8, 7 + header_len).decode('utf-8'))
except Exception as e:
return {"metadata": None, "error": str(e)}
from modules import model_probe
response = {
"metadata": header.get('__metadata__'),
"tensors": len([k for k in header if k != '__metadata__']),
"probe": model_probe.analyze_header(header),
}
if file_id:
peek_cache_put(file_id, url, response)
return response
# ---------------------------------------------------------------------------
# Settings
# ---------------------------------------------------------------------------
@@ -585,6 +695,7 @@ def register_api(api):
api.add_api_route("/sdapi/v2/civitai/download", post_download, methods=["POST"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/download/{download_id}/cancel", post_download_cancel, methods=["POST"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/download/status", get_download_status, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/peek-header", get_peek_header, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/settings", get_settings, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/settings", post_settings, methods=["POST"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/resolve-path", get_resolve_path, methods=["GET"], tags=["CivitAI"])
modules/civitai/download_civitai.py
+35 -1
View File
@@ -1,4 +1,5 @@
import os
import re
import uuid
import hashlib
import threading
@@ -12,6 +13,10 @@ from modules import shared, paths
from modules.logger import console
STALE_PARTIAL_DAYS = 7
TEMP_FILE_RE = re.compile(r'^[0-9a-f]{8}\.tmp$')
@dataclass
class DownloadItem:
id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
@@ -57,6 +62,27 @@ class DownloadManager:
self._lock = threading.Lock()
self._max_workers = max_workers
self._worker_count = 0
threading.Thread(target=self._sweep_stale_partials, daemon=True).start()
def _sweep_stale_partials(self):
# Partials double as resume state, but mtime stops moving the moment a
# download stops, so an untouched week-old .tmp is abandoned.
try:
from modules.civitai.filemanage_civitai import iter_type_roots
cutoff = time.time() - STALE_PARTIAL_DAYS * 86400
for root in iter_type_roots():
for path in root.rglob('*.tmp'):
if not TEMP_FILE_RE.match(path.name):
continue
try:
stat = path.stat()
if stat.st_mtime < cutoff:
path.unlink()
log.info(f'CivitAI stale partial removed: file="{path}" size={stat.st_size/1024/1024:.0f}MB age>{STALE_PARTIAL_DAYS}d')
except OSError as e:
log.warning(f'CivitAI stale partial sweep: file="{path}" {e}')
except Exception as e:
log.warning(f'CivitAI stale partial sweep error: {e}')
def enqueue(self, url: str, folder: str, filename: str, model_type: str = "",
expected_hash: str = "", token: str | None = None,
@@ -188,6 +214,15 @@ class DownloadManager:
log.warning(f'CivitAI download invalid content-type: id={item.id} content-type="{content_type}"')
return
# A 200 reply to a Range request means the server ignored the range
# and is sending the whole file; appending it to the partial would
# corrupt it. Truncate and restart from byte 0.
if starting_pos > 0 and r.status_code == 200:
log.warning(f'CivitAI download resume not supported: id={item.id} file="{item.filename}" restarting')
starting_pos = 0
item.bytes_downloaded = 0
os.truncate(temp_file, 0)
total_size = int(r.headers.get('content-length', 0))
item.bytes_total = starting_pos + total_size
@@ -742,7 +777,6 @@ def download_civit_model(model_url: str, model_name: str = '', model_path: str =
log.error('Model download: no url provided')
return None
if not version_id:
import re
match = re.search(r'/api/download/models/(\d+)', model_url)
if match:
version_id = int(match.group(1))
modules/civitai/filemanage_civitai.py
+23 -1
View File
@@ -4,9 +4,12 @@ from pathlib import Path
from modules.logger import log
# Map CivitAI model types to shared.opts directory settings and fallback subfolder names
# Map CivitAI model types to shared.opts directory settings and fallback subfolder
# names. 'Text Encoder' is a file type, not a model type: versions bundle companion
# files, and clients route those by the file's own type.
TYPE_MAP = {
'Checkpoint': ('ckpt_dir', 'Stable-diffusion'),
'Text Encoder': ('te_dir', 'Text-encoder'),
'TextualInversion': ('embeddings_dir', 'embeddings'),
'Hypernetwork': ('hypernetwork_dir', 'hypernetworks'),
'AestheticGradient': ('ckpt_dir', 'Stable-diffusion'),
@@ -60,6 +63,25 @@ def get_type_folder(model_type: str, base_model: str = '') -> Path:
return Path(paths.models_path) / fallback_dir
def iter_type_roots() -> set[Path]:
"""Every root folder downloads can resolve into, for maintenance sweeps."""
from modules import shared, paths
roots = set()
custom_json = getattr(shared.opts, 'civitai_save_type_folders', '') or ''
if custom_json.strip():
try:
import json
for folder in json.loads(custom_json).values():
p = Path(folder)
roots.add(p if p.is_absolute() else Path(paths.models_path) / folder)
except Exception:
pass
for opt_attr, fallback_dir in set(TYPE_MAP.values()) | {('unet_dir', 'UNET')}:
configured = (getattr(shared.opts, opt_attr, '') or '') if opt_attr else ''
roots.add(Path(configured) if configured else Path(paths.models_path) / fallback_dir)
return {r for r in roots if r.is_dir()}
def resolve_save_path(model_type: str, model_name: str = "", base_model: str = "",
nsfw: bool = False, creator: str = "", model_id: int = 0,
version_id: int = 0, version_name: str = "") -> Path:
modules/civitai/models_civitai.py
+22
View File
@@ -25,6 +25,15 @@ class CivitFileHashes(BaseModel):
blake3: str | None = Field(None, alias="BLAKE3")
class CivitFileMetadata(BaseModel):
class Config:
allow_population_by_field_name = True
format: str | None = None
size: str | None = None
fp: str | None = None
quant_type: str | None = Field(None, alias="quantType")
class CivitFile(BaseModel):
class Config:
allow_population_by_field_name = True
@@ -35,6 +44,10 @@ class CivitFile(BaseModel):
hashes: CivitFileHashes = Field(default_factory=CivitFileHashes)
download_url: str = Field("", alias="downloadUrl")
primary: bool | None = None
metadata: CivitFileMetadata = Field(default_factory=CivitFileMetadata)
pickle_scan_result: str | None = Field(None, alias="pickleScanResult")
virus_scan_result: str | None = Field(None, alias="virusScanResult")
scanned_at: str | None = Field(None, alias="scannedAt")
class CivitStats(BaseModel):
@@ -58,6 +71,8 @@ class CivitVersion(BaseModel):
base_model: str = Field("Unknown", alias="baseModel")
published_at: str | None = Field(None, alias="publishedAt")
availability: str = "Unknown"
early_access_ends_at: str | None = Field(None, alias="earlyAccessEndsAt")
early_access_config: dict | None = Field(None, alias="earlyAccessConfig")
description: str | None = None
trained_words: list[str] = Field(default_factory=list, alias="trainedWords")
stats: CivitStats = Field(default_factory=CivitStats)
@@ -66,6 +81,13 @@ class CivitVersion(BaseModel):
nsfw_level: int = Field(0, alias="nsfwLevel")
download_url: str = Field("", alias="downloadUrl")
@validator('availability', pre=True)
def coerce_null_availability(cls, v): # pylint: disable=no-self-argument
# /model-versions/{id} serializes availability as null, which failed
# str validation and turned every version lookup into a 404. Fall back
# to the default instead of rejecting the whole version.
return "Unknown" if v in (None, "") else v
class CivitCreator(BaseModel):
class Config:
modules/civitai/userdata_civitai.py
+3 -1
View File
@@ -96,13 +96,15 @@ class SearchHistory:
except Exception as e:
log.error(f'CivitAI search history save error: file={path} {e}')
def add(self, search_type: str, term: str):
def add(self, search_type: str, term: str, params: dict | None = None):
with self._lock:
entry = {
"type": search_type,
"term": term,
"timestamp": datetime.now().isoformat(),
}
if params:
entry["params"] = params
# Remove duplicate if same type+term exists
self._entries = [e for e in self._entries if not (e.get('type') == search_type and e.get('term') == term)]
self._entries.insert(0, entry)
modules/infotext.py
+4 -6
View File
@@ -66,12 +66,10 @@ def parse(infotext):
return params
params['Prompt'] = prompt if len(prompt) > 0 else None
params['Negative prompt'] = negative if len(negative) > 0 else None
params['Template'] = template if len(template) > 0 else None
params['Negative template'] = negative_template if len(negative_template) > 0 else None
if params['Template'] == params['Prompt']:
params.pop('Template', None)
if params['Negative template'] == params['Negative prompt']:
params.pop('Negative template', None)
if len(template) > 0 and template != params['Prompt']:
params['Template'] = template
if len(negative_template) > 0 and negative_template != params['Negative prompt']:
params['Negative template'] = negative_template
debug(f'Params: {params}')
for key, val in params.copy().items():
modules/lora/native_adapter.py
+208 -18
View File
@@ -48,7 +48,11 @@ from modules.lora import lora_common as l
# vendor-specific naming conventions. ``lora_transformer_`` is not a vendor
# format but sdnext's own internal transformer namespace; files already saved
# in it (e.g. OneTrainer) pass through verbatim, see :func:`resolve_group_targets`.
KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_", "lora_transformer_")
# ``lycoris_`` is the LyCORIS-standalone save format: the wrapped diffusers
# module path with dots rendered as underscores. It is arch-independent by
# construction (the wrapped layout equals the target layout), so it lives here
# rather than in any arch's prefix list.
KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_", "lora_transformer_", "lycoris_")
# Sentinel ``prefix_used`` value emitted by :func:`parse_key` when a bare path
@@ -67,9 +71,9 @@ NETWORK_PREFIX_DEFAULT = "lora_transformer_"
# Prefixes whose parsed ``base`` is already a network-key tail (``arch_prefix +
# base.replace(".", "_")`` matches the stamped module name), so the loader binds
# them directly with no per-arch rewrite. Arch-local already-resolved prefixes
# (e.g. flux2's ``lycoris_``) stay in that arch's ``resolve_targets``.
PASSTHROUGH_PREFIXES_DEFAULT = ("transformer.", BARE_DIFFUSERS_PREFIX_USED, "lora_transformer_")
# them directly with no per-arch rewrite. ``lycoris_`` bases are already
# underscored, so the loader's ``.replace(".", "_")`` is a no-op on them.
PASSTHROUGH_PREFIXES_DEFAULT = ("transformer.", BARE_DIFFUSERS_PREFIX_USED, "lora_transformer_", "lycoris_")
def _resolve_prefix(network_prefix, prefix_used):
@@ -98,7 +102,12 @@ LORA_SUFFIXES = (
".lora_down.weight", ".lora_up.weight", ".lora_mid.weight",
".lora_A.weight", ".lora_B.weight",
# diff_b: bias delta some saves pair with the weight LoRA, applied as ex_bias.
".alpha", ".dora_scale", ".bias", ".diff_b", ".scale",
# magnitude / lora_magnitude_vector: DoRA row norms (ai-toolkit / PEFT key
# names); converted onto the dora_scale path by try_load_lora.
# bias_indices/values/size: LyCORIS extraction sparse residual triplet,
# reconstructed into the dense-bias path by network.NetworkModule.
".alpha", ".dora_scale", ".magnitude", ".lora_magnitude_vector",
".bias", ".bias_indices", ".bias_values", ".bias_size", ".diff_b", ".scale",
)
LOKR_SUFFIXES = (
".lokr_w1", ".lokr_w2",
@@ -288,6 +297,63 @@ def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0]
def lokr_kron_shape(w):
"""Return the ``(out, in_flat)`` shape of ``kron(w1, w2)`` from stored factors.
Each factor is either full (``lokr_w1``/``lokr_w2``) or rank-decomposed
(``lokr_w1_a @ lokr_w1_b``); a Tucker-rebuilt w2 (conv-only) stores its
parts as ``(rank, part)`` with the kernel dims carried by ``lokr_t2``.
Kernel dims are folded into ``in_flat``, matching how the delta reshapes
onto a conv weight.
"""
w1 = w.get("lokr_w1")
r1 = w1.shape[0] if w1 is not None else w["lokr_w1_a"].shape[0]
c1 = w1.shape[1] if w1 is not None else w["lokr_w1_b"].shape[1]
w2 = w.get("lokr_w2")
t2 = w.get("lokr_t2")
if w2 is not None:
r2 = w2.shape[0]
c2_flat = w2.numel() // r2
elif t2 is not None:
r2 = w["lokr_w2_a"].shape[1]
c2_flat = w["lokr_w2_b"].shape[1]
for d in t2.shape[2:]:
c2_flat *= d
else:
r2 = w["lokr_w2_a"].shape[0]
w2b = w["lokr_w2_b"]
c2_flat = w2b.numel() // w2b.shape[0]
return r1 * r2, c1 * c2_flat
def lokr_shapes_match(sd_module, kron_shape, chunk: ChunkSpec | None) -> bool:
"""Kron-vs-module dim check, honoring SDNQ original shapes and chunk rows.
The input dim is never chunked; the output dim must cover the full fused
weight for chunked targets (``total * out`` for equal chunks, ``end <=
kron_out`` with an exact row-range for slices).
"""
if not hasattr(sd_module, "weight"):
return False
if hasattr(sd_module, "sdnq_dequantizer"):
mod_shape = sd_module.sdnq_dequantizer.original_shape
else:
mod_shape = sd_module.weight.shape
if len(mod_shape) < 2:
return False
mod_in_flat = 1
for d in mod_shape[1:]:
mod_in_flat *= d
kron_out, kron_in_flat = kron_shape
if kron_in_flat != mod_in_flat:
return False
if chunk is None:
return kron_out == mod_shape[0]
if chunk.is_equal_chunks:
return kron_out == mod_shape[0] * chunk.total
return (chunk.end - chunk.start) == mod_shape[0] and kron_out >= chunk.end
# === Parsing primitives ===
@@ -403,27 +469,75 @@ def resolve_group_targets(resolve_targets, prefix_used, base):
# slices and Tucker-decomposed LoHAs on fused targets are skipped with a
# warning (no slice variant exists, and Tucker keys cannot arise on Linear
# layers per LyCORIS upstream — see network_hada.NetworkModuleHadaChunk).
#
# DoRA on fused targets (LoRA/LoKR/LoHA): per-output dora_scale rows are
# sliced with the chunk via slice_dora_scale (exact, row norms are
# independent); per-input dora_scale couples the chunks through shared
# column norms and the group is skipped with a warning.
#
# Bias companions on fused targets: per-output diff_b deltas slice with the
# chunk (LoRA only; no other family's save format pairs them). The legacy
# weight-shaped "bias" key has no defined partition and skips the group.
# - OFT/BOFT: no chunk variant exists; fused targets are skipped with a
# warning. Discrimination is by ``oft_blocks.ndim`` (3-D OFT, 4-D BOFT),
# mirroring upstream LyCORIS ``algo_check``.
def _slice_lora_chunk(w, chunk: ChunkSpec):
"""Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``.
def slice_chunk_rows(t, chunk: ChunkSpec):
"""Slice dim 0 of ``t`` per ``chunk``.
Equal-chunks form uses ``torch.chunk`` (faster for the symmetric case);
row-range form uses tensor slicing for arbitrary partitions.
"""
up = w["lora_up.weight"]
if chunk.is_equal_chunks:
sliced = torch.chunk(up, chunk.total, dim=0)[chunk.idx].contiguous()
else:
sliced = up[chunk.start:chunk.end].contiguous()
return torch.chunk(t, chunk.total, dim=0)[chunk.idx].contiguous()
return t[chunk.start:chunk.end].contiguous()
def _slice_lora_chunk(w, chunk: ChunkSpec):
"""Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``."""
out = dict(w)
out["lora_up.weight"] = sliced
out["lora_up.weight"] = slice_chunk_rows(w["lora_up.weight"], chunk)
return out
def slice_dora_scale(w, chunk: ChunkSpec, fused_out):
"""Slice a per-output ``dora_scale`` to the chunk rows; ``None`` if unsliceable.
LyCORIS ``wd_on_out=True`` (the LoKr/LoHA default) stores per-output
magnitudes of shape ``(out, 1)`` (1-D ``(out,)`` also seen); the rows
partition exactly with the fused weight, so the chunk slice preserves the
DoRA math (row norms are independent across rows). ``wd_on_out=False``
stores per-input magnitudes whose column norms span every fused row,
coupling the chunks; no exact per-chunk equivalent exists and the caller
must skip the group.
"""
ds = w.get("dora_scale")
if ds is None:
return w
if ds.ndim >= 1 and ds.shape[0] == fused_out:
out = dict(w)
out["dora_scale"] = slice_chunk_rows(ds, chunk)
return out
return None
def slice_bias_delta(w, chunk: ChunkSpec, fused_out):
"""Slice a per-output ``diff_b`` bias delta to the chunk rows; ``None`` if unsliceable.
``diff_b`` stores one bias value per output feature, so it partitions with
the fused rows exactly like the up-weight.
"""
db = w.get("diff_b")
if db is None:
return w
if db.ndim >= 1 and db.shape[0] == fused_out:
out = dict(w)
out["diff_b"] = slice_chunk_rows(db, chunk)
return out
return None
def try_load_lora(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
@@ -450,9 +564,20 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
unmapped = 0
mismatch = 0
skipped = 0
for (prefix, base), w in groups.items():
if "lora_down.weight" not in w or "lora_up.weight" not in w:
continue
# DoRA magnitude vectors: ai-toolkit saves `magnitude`, PEFT/diffusers
# `lora_magnitude_vector`. Both are 1-D per-output row norms with
# dora_scale semantics; reshape to (out, 1) so the apply-time
# orientation detection cannot misread square layers as per-input.
for mag_key in ("magnitude", "lora_magnitude_vector"):
mag = w.get(mag_key)
if mag is not None and "dora_scale" not in w:
w = dict(w)
w.pop(mag_key)
w["dora_scale"] = mag.reshape(-1, 1) if mag.ndim == 1 else mag
arch_prefix = _resolve_prefix(network_prefix, prefix)
for diffusers_path, chunk in resolve_group_targets(resolve_targets, prefix, base):
network_key = arch_prefix + diffusers_path.replace(".", "_")
@@ -461,7 +586,26 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
unmapped += 1
continue
target_w = _slice_lora_chunk(w, chunk) if chunk is not None else w
target_w = w
if chunk is not None:
if "bias" in w or "bias_indices" in w:
# Weight-shaped bias residuals (dense or LyCORIS sparse
# triplet) are not partitioned onto fused targets.
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
continue
fused_out = w["lora_up.weight"].shape[0]
target_w = _slice_lora_chunk(w, chunk)
target_w = slice_dora_scale(target_w, chunk, fused_out)
if target_w is None:
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)')
skipped += 1
continue
target_w = slice_bias_delta(target_w, chunk, fused_out)
if target_w is None:
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} non-per-output diff_b on fused target skipped (unsupported)')
skipped += 1
continue
if not shapes_match(sd_module, target_w["lora_down.weight"], target_w["lora_up.weight"]):
log.warning(
@@ -476,14 +620,14 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module)
net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
return finalize_network(net, name, "LoRA", lora_scale, t0, unmapped=unmapped, mismatch=mismatch)
return finalize_network(net, name, "LoRA", lora_scale, t0, unmapped=unmapped, mismatch=mismatch, skipped=skipped)
def try_load_lokr(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
arch_name="generic"): # pylint: disable=unused-argument
arch_name="generic"):
"""Generic LoKR loader.
Stores only the compact LoKR factors and dispatches to
@@ -508,19 +652,48 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
)
unmapped = 0
mismatch = 0
skipped = 0
for (prefix, base), w in groups.items():
has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w)
has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w)
if not (has_1 and has_2):
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
kron_shape = lokr_kron_shape(w)
for diffusers_path, chunk in resolve_group_targets(resolve_targets, prefix, base):
network_key = arch_prefix + diffusers_path.replace(".", "_")
sd_module = mapping.get(network_key)
if sd_module is None:
unmapped += 1
continue
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
if not lokr_shapes_match(sd_module, kron_shape, chunk):
log.warning(
f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key}'
f' kron={kron_shape[0]}x{kron_shape[1]}'
f' module={getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}'
f' shape mismatch'
)
mismatch += 1
continue
target_w = w
if chunk is not None:
if "bias" in w:
log.warning(f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
continue
# Kron rows = w1 rows * w2 rows; the tucker w2_a orientation
# differs but tucker is conv-only and chunks are Linear-only.
w1 = w.get("lokr_w1")
w2 = w.get("lokr_w2")
w1_rows = w1.shape[0] if w1 is not None else w["lokr_w1_a"].shape[0]
w2_rows = w2.shape[0] if w2 is not None else w["lokr_w2_a"].shape[0]
target_w = slice_dora_scale(w, chunk, fused_out=w1_rows * w2_rows)
if target_w is None:
log.warning(f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)')
skipped += 1
continue
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module)
if chunk is None:
net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw)
elif chunk.is_equal_chunks:
@@ -528,7 +701,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
else:
net.modules[network_key] = network_lokr.NetworkModuleLokrSliceChunk(net, nw, chunk.start, chunk.end)
return finalize_network(net, name, "LoKR", lora_scale, t0, unmapped=unmapped)
return finalize_network(net, name, "LoKR", lora_scale, t0, unmapped=unmapped, mismatch=mismatch, skipped=skipped)
def try_load_loha(name, network_on_disk, lora_scale, *,
@@ -577,7 +750,18 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
if sd_module is None:
unmapped += 1
continue
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
target_w = w
if chunk is not None:
if "bias" in w:
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
continue
target_w = slice_dora_scale(w, chunk, fused_out=w["hada_w1_a"].shape[0])
if target_w is None:
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)')
skipped += 1
continue
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module)
if chunk is None:
net.modules[network_key] = network_hada.NetworkModuleHada(net, nw)
elif chunk.is_equal_chunks:
@@ -855,6 +1039,12 @@ def try_load_full(name, network_on_disk, lora_scale, *,
if sd_module is None:
unmapped += 1
continue
# Loader-local stamping, same as try_load_norm: a full-weight extraction carries the
# norm weights too, and assign_network_names_to_compvis_modules puts norms in the
# mapping but never stamps network_layer_name on them, which is what the apply pass
# keys off. Without this they bind and then silently never apply.
if not getattr(sd_module, "network_layer_name", None):
sd_module.network_layer_name = network_key
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
net.modules[network_key] = network_full.NetworkModuleFull(net, nw)
modules/lora/network.py
+18 -1
View File
@@ -1,6 +1,7 @@
import os
import enum
from collections import namedtuple
import torch
from modules import hashes, shared, sd_checkpoint
@@ -174,6 +175,16 @@ class NetworkModule:
self.shape = self.sd_module.weight.shape
self.dim = None
self.bias = weights.w.get("bias")
if self.bias is None and "bias_indices" in weights.w:
# LyCORIS extraction with use_bias: the sparse weight-shaped
# remainder of the SVD extraction ("bias" is historical naming),
# stored COO with int16 indices. Kept sparse; finalize_updown's
# dense += sparse materializes it per module at apply time.
self.bias = torch.sparse_coo_tensor(
weights.w["bias_indices"].to(torch.long),
weights.w["bias_values"],
tuple(weights.w["bias_size"]),
)
self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None
self.scale = weights.w["scale"].item() if "scale" in weights.w else None
self.dora_scale = weights.w.get("dora_scale", None)
@@ -268,7 +279,13 @@ class NetworkModule:
if ex_bias is not None:
ex_bias = ex_bias * self.multiplier()
if self.dora_scale is not None:
updown = self.apply_weight_decompose(updown, orig_weight)
# LyCORIS/ComfyUI convention: alpha/rank is baked into the diff
# before the decompose norm. The multiplier then lerps the full
# merged delta (ComfyUI semantics: 0 disables, 1 equals the
# trainer's output; LyCORIS weight-mode ratio interpolation is
# not used since it leaves the diff applied at multiplier 0).
updown = self.apply_weight_decompose(updown * self.calc_scale(), orig_weight)
return updown * self.multiplier(), ex_bias
return updown * self.calc_scale() * self.multiplier(), ex_bias
def calc_updown(self, target):
modules/ltx/ltx_ui.py
+8 -1
View File
@@ -4,7 +4,7 @@ from modules import ui_sections, ui_symbols
from modules.ui_components import ToolButton
from modules.logger import log
from modules.video_models.models_def import models
from modules.ltx import ltx_process, ltx_capabilities
from modules.ltx import ltx_process, ltx_capabilities, ltx_util
debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -63,6 +63,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
with gr.Row():
ltx_models = [m.name for m in models['LTX Video']] if 'LTX Video' in models else ['None']
model = gr.Dropdown(label='LTX model', choices=ltx_models, value=ltx_models[0], elem_id="ltx_model")
btn_load = ToolButton(ui_symbols.loading, elem_id="video_model_load_ltx")
with gr.Accordion(open=False, label='Size', elem_id='ltx_size_accordion'):
width, height = ui_sections.create_resolution_inputs('ltx', default_width=832, default_height=480)
with gr.Row():
@@ -119,6 +120,12 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
with gr.Row():
text = gr.HTML('', elem_id='ltx_generation_info', show_label=False)
def load_model(model_name: str):
ltx_util.load_model('LTX Video', model_name)
btn_load.click(fn=load_model, inputs=[model], outputs=[])
model.change(
fn=_model_change,
inputs=[model],
modules/ltx/ltx_util.py
+3 -3
View File
@@ -28,7 +28,7 @@ def load_model(engine: str, model: str):
selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0]
# video_load owns the cache; pipe-class mismatch inside it invalidates the name-based hit
# when Unload Models (or any external swap) silently replaced shared.sd_model.
log.info(f'Video load: engine="{engine}" selected="{model}" {selected}')
log.info(f'Load video: engine="{engine}" selected="{model}" {selected}')
video_load.load_model(selected)
t1 = time.time()
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
@@ -41,7 +41,7 @@ def load_upsample(upsample_pipe, upsample_repo_id):
if upsample_pipe is None:
t0 = time.time()
from diffusers.pipelines.ltx.pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline
log.info(f'Video load: cls={LTXLatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
log.info(f'Load video: cls={LTXLatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
upsample_pipe = LTXLatentUpsamplePipeline.from_pretrained(
upsample_repo_id,
vae=shared.sd_model.vae,
@@ -61,7 +61,7 @@ def load_upsample_2x(upsample_pipe, upsample_repo_id):
from diffusers.pipelines.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from modules import sd_checkpoint
log.info(f'Video load: cls={LTX2LatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
log.info(f'Load video: cls={LTX2LatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
upsample_repo_id,
subfolder='latent_upsampler',
modules/model_probe.py
+429
View File
@@ -0,0 +1,429 @@
"""Header-only model analysis: architecture fingerprinting, precision/quant
detection, and embedded metadata extraction from safetensors/gguf files
without reading tensor data. Consumed by the civitai remote peek and the
local model audit; analyze_header is pure so both paths share one analyzer.
"""
import os
import re
import json
from collections import Counter
from dataclasses import dataclass
SCHEMA_VERSION = 5
MAX_HEADER_BYTES = 16 * 1024 * 1024
STRIP_PREFIXES = ('model.diffusion_model.', 'diffusion_model.', 'net.')
# companion families bundled alongside the diffusion core in full checkpoints
COMPANION_PREFIXES = ('text_encoders.', 'cond_stage_model.', 'conditioner.', 'first_stage_model.', 'vae.', 'vocoder.', 'audio_vae.')
LORA_SUFFIXES = ('.lora_down.weight', '.lora_up.weight', '.lora_A.weight', '.lora_B.weight', '.hada_w1_a', '.lokr_w1', '.lokr_w2', '.dora_scale', '.diff', '.diff_b')
@dataclass(frozen=True)
class ArchFingerprint:
family: str
display: str
required: tuple # regex fragments, all must match at least one inner key
forbidden: tuple = ()
weight: int = 1
# Markers verified against real headers (local library + civitai ranged peeks).
# Families not listed report 'unknown' rather than guessing.
FINGERPRINTS = (
ArchFingerprint('sdxl', 'Stable Diffusion XL', (r'^input_blocks\.', r'^middle_block\.', r'^label_emb\.')),
ArchFingerprint('sd', 'Stable Diffusion 1.x/2.x', (r'^input_blocks\.', r'^middle_block\.'), forbidden=(r'^label_emb\.',)),
ArchFingerprint('sd3', 'Stable Diffusion 3', (r'^joint_blocks\.', r'^context_embedder\.')),
ArchFingerprint('f1', 'FLUX.1', (r'^double_blocks\.\d+\.img_attn', r'^vector_in\.'), forbidden=(r'^distilled_guidance_layer\.', r'stream_modulation')),
ArchFingerprint('chroma', 'Chroma', (r'^double_blocks\.\d+\.img_attn', r'^distilled_guidance_layer\.')),
ArchFingerprint('f2', 'FLUX.2', (r'^double_blocks\.\d+\.img_attn', r'stream_modulation'), forbidden=(r'^vector_in\.',)),
ArchFingerprint('qwen', 'Qwen Image', (r'^transformer_blocks\.\d+\.img_mlp', r'^time_text_embed\.'), forbidden=(r'audio_attn',)),
ArchFingerprint('krea2', 'Krea 2', (r'^blocks\.\d+\.attn', r'^txtfusion\.')),
ArchFingerprint('wanai-animate', 'Wan 2.2 Animate', (r'^blocks\.\d+\.cross_attn', r'^patch_embedding\.', r'^face_adapter\.|^motion_encoder\.|^pose_patch_embedding\.')),
ArchFingerprint('wanai', 'Wan DiT', (r'^blocks\.\d+\.cross_attn', r'^blocks\.\d+\.self_attn', r'^patch_embedding\.', r'^time_projection\.'), forbidden=(r'^face_adapter\.', r'^motion_encoder\.', r'^pose_patch_embedding\.', r'^llm_adapter\.')),
ArchFingerprint('anima', 'Anima', (r'^blocks\.\d+\.cross_attn', r'^llm_adapter\.')),
ArchFingerprint('ernieimage', 'Ernie Image', (r'^layers\.\d+\.self_attention', r'^adaLN_sa_ln|^final_norm\.')),
ArchFingerprint('zimage', 'Z-Image', (r'^layers\.\d+\.attention', r'^noise_refiner\.', r'^cap_embedder\.')),
ArchFingerprint('ltxvideo', 'LTX Video', (r'^transformer_blocks\.\d+\.attn1', r'^adaln_single\.|^vocoder\.')),
ArchFingerprint('ideogram4', 'Ideogram 4', (r'^layers\.\d+\.attention\.(qkv|o)\.', r'^adaln_proj\.|^t_embedding\.')),
)
# name/text to family hints, ordered specific-first; matched against lowercased
# text with spaces removed. Used for lora trainer metadata and for names that
# imply a base (folder names, base tags).
FAMILY_NAME_HINTS = (
('stable-diffusion-xl', 'sdxl'), ('sdxl', 'sdxl'), ('sd_xl', 'sdxl'),
('animagine', 'sdxl'), ('pony', 'sdxl'), ('illustrious', 'sdxl'), ('noobai', 'sdxl'),
('stable-diffusion-3', 'sd3'), ('sd3', 'sd3'),
('stable-diffusion-v1', 'sd'), ('sd-v1', 'sd'), ('sd_v1', 'sd'), ('sd15', 'sd'), ('sd1.5', 'sd'), ('sd2', 'sd'),
('flux.2', 'f2'), ('flux2', 'f2'), ('klein', 'f2'), ('chroma', 'chroma'), ('flux', 'f1'),
('krea', 'krea2'), ('qwen', 'qwen'), ('wan', 'wanai'), ('anima', 'anima'),
('ideogram', 'ideogram4'), ('z-image', 'zimage'), ('zimage', 'zimage'),
('ernie', 'ernieimage'), ('ltx', 'ltxvideo'),
)
def family_from_text(text: str) -> str | None:
compact = (text or '').lower().replace(' ', '')
for hint, family in FAMILY_NAME_HINTS:
if hint in compact:
return family
return None
LORA_KEY_HINTS = (
('txtfusion', 'krea2'),
('distilled_guidance', 'chroma'),
('stream_modulation', 'f2'),
('double_blocks', 'f1'),
('joint_blocks', 'sd3'),
('llm_adapter', 'anima'),
('cross_attn', 'wanai'),
('img_mlp', 'qwen'),
('noise_refiner', 'zimage'),
('context_refiner', 'zimage'),
# diffusers flux layout; f1 vs f2 not separable from lora keys
('single_transformer_blocks', 'f1'),
('transformer_blocks', 'qwen'),
)
# checked in order; trainers sometimes write a bogus modelspec.architecture
# while ss_base_model_version names the real base
LORA_META_KEYS = ('ss_base_model_version', 'modelspec.architecture', 'ss_sd_model_name')
FAMILY_DISPLAY = {fp.family: fp.display for fp in FINGERPRINTS}
# expert-role vocabulary for dual-transformer archs, matched against
# __metadata__ values; wan uses phrases because a bare 'high' appears in
# unrelated metadata (e.g. 'high quality')
ROLE_WORDS = {
'wanai': (('high-noise', ('high noise',)), ('low-noise', ('low noise',))),
'ideogram4': (('uncond', ('uncond', 'unconditional')),),
}
def strip_key(key: str) -> str:
for prefix in STRIP_PREFIXES:
if key.startswith(prefix):
return key[len(prefix):]
return key
def split_words(text: str) -> str:
text = re.sub(r'([a-z0-9])([A-Z])', r'\1 \2', text or '')
return re.sub(r'[^a-zA-Z0-9]+', ' ', text).lower()
def detect_kind(inner_keys: list) -> str:
"""Container-level classification before arch matching."""
if any(k.endswith(LORA_SUFFIXES) or '.lora_down.' in k or '.lora_up.' in k or '.lora.down.' in k or '.lora.up.' in k or '.lora_A.' in k or '.lora_B.' in k for k in inner_keys):
return 'lora'
top = Counter(k.split('.')[0] for k in inner_keys)
total = sum(top.values())
vae_keys = top.get('encoder', 0) + top.get('decoder', 0) + top.get('quant_conv', 0) + top.get('post_quant_conv', 0)
if total > 0 and top.get('decoder', 0) > 0 and vae_keys > total * 0.9:
return 'vae'
diffusion_core = any(k.startswith(('input_blocks', 'double_blocks', 'blocks.', 'layers.', 'transformer_blocks', 'joint_blocks')) for k in inner_keys)
llm_markers = any(('.mlp.gate_proj' in k or 'embed_tokens' in k or '.DenseReluDense.' in k) for k in inner_keys) or top.get('shared', 0) > 0
if not diffusion_core and llm_markers:
return 'text-encoder'
return 'model'
def lora_context_dim(shapes: dict | None) -> int | None:
"""Cross-attention context width from a lora down/A tensor: 768 for sd1.x,
1024 for sd2.x, 2048 for sdxl. Arbitrates when trainer metadata lies."""
for k, shape in (shapes or {}).items():
if 'attn2' in k and 'to_k' in k and ('lora_down' in k or '.lora.down.' in k or 'lora_A' in k):
if isinstance(shape, (list, tuple)) and len(shape) == 2:
return int(shape[1])
return None
def match_lora_base(core_keys: list, metadata: dict | None, shapes: dict | None = None):
"""Loras train a module subset, so full fingerprints cannot apply. Trainer
metadata names the base; keyless PEFT loras resolve via distinctive
substrings. Returns (family, confidence, marker) or None."""
def unet_split(family, confidence, marker):
# sd-vs-sdxl from metadata or key heuristics is unreliable; shapes win
dim = lora_context_dim(shapes)
if dim is not None:
return ('sdxl' if dim >= 2048 else 'sd'), 0.8, f'shape:context-dim={dim}'
return family, confidence, marker
for key in LORA_META_KEYS:
family = family_from_text(str((metadata or {}).get(key, '')))
if family is not None:
if family in ('sd', 'sdxl'):
return unet_split(family, 0.9, f'metadata:{key}')
return family, 0.9, f'metadata:{key}'
joined = '\n'.join(core_keys)
for hint, family in LORA_KEY_HINTS:
if hint in joined:
return family, 0.5, f'key:{hint}'
# bare blocks.N.attn without self/cross variants is the krea2 DiT layout
if re.search(r'blocks\.\d+\.attn\.', joined) and 'self_attn' not in joined:
return 'krea2', 0.5, 'key:blocks.attn'
if 'input_blocks' in joined or 'down_blocks' in joined:
return unet_split('sdxl' if 'te2' in joined else 'sd', 0.5, 'key:unet-blocks')
return None
def match_arch(keys: list, metadata: dict | None, shapes: dict | None = None) -> dict:
inner = [strip_key(k) for k in keys]
core = [k for k in inner if not k.startswith(COMPANION_PREFIXES)]
prefixes = {k[: len(k) - len(s)] for k, s in ((k, strip_key(k)) for k in keys) if k != s}
kind = detect_kind(core)
matched_family = None
matched_markers = []
candidates = []
confidence = 0.0
if kind == 'model':
for fp in FINGERPRINTS:
hits = [pat for pat in fp.required if any(re.search(pat, k) for k in core)]
if len(hits) < len(fp.required):
continue
if any(any(re.search(pat, k) for k in core) for pat in fp.forbidden):
continue
candidates.append(fp)
if matched_family is None:
matched_family = fp.family
matched_markers = hits
confidence = 1.0 if matched_family and len(candidates) == 1 else (0.7 if matched_family else 0.0)
elif kind == 'lora':
resolved = match_lora_base(core, metadata, shapes)
if resolved is not None:
matched_family, confidence, marker = resolved
matched_markers = [marker]
variant = None
if matched_family is not None and metadata:
meta_text = split_words(' '.join(str(v) for v in metadata.values() if isinstance(v, str)))
for suffix, words in ROLE_WORDS.get(matched_family, ()):
if any(re.search(rf'\b{w}\b', meta_text) for w in words):
variant = suffix
break
return {
'kind': kind,
'family': matched_family or 'unknown',
'display': FAMILY_DISPLAY.get(matched_family, 'Unknown') if matched_family else 'Unknown',
'confidence': confidence,
'variant': variant,
'detected_prefix': sorted(prefixes)[0] if prefixes else '',
'matched_markers': matched_markers,
'candidates': [{'family': fp.family, 'display': fp.display} for fp in candidates],
}
# storage dtype of a marked layer's weight identifies its comfy_quant format
COMFY_QUANT_DTYPE_FORMATS = {'I8': 'int8_tensorwise', 'F8_E4M3': 'float8_e4m3fn', 'U8': 'nvfp4'}
def detect_quant(keys: list, dtypes: Counter, container: str, metadata: dict | None = None, key_dtypes: dict | None = None) -> dict:
if container == 'gguf':
quantized = {d: n for d, n in dtypes.items() if d not in ('F32', 'F16', 'BF16')}
dominant = max(quantized, key=quantized.get) if quantized else None
return {'scheme': 'gguf' if dominant else None, 'format': dominant, 'marked_layers': None, 'source': 'gguf-qtype'}
quant_metadata = (metadata or {}).get('_quantization_metadata')
if quant_metadata:
try:
layers = (json.loads(quant_metadata) if isinstance(quant_metadata, str) else quant_metadata).get('layers') or {}
except Exception:
layers = {}
formats = Counter(entry.get('format') for entry in layers.values() if isinstance(entry, dict) and entry.get('format'))
if formats:
return {'scheme': 'comfy_quant', 'format': formats.most_common(1)[0][0], 'marked_layers': len(layers), 'source': 'header'}
comfy = [k for k in keys if k.endswith('.comfy_quant')]
if comfy:
weight_dtypes = Counter()
for marker in comfy:
dtype = (key_dtypes or {}).get(f"{marker[: -len('.comfy_quant')]}.weight")
if dtype:
weight_dtypes[dtype] += 1
dominant = weight_dtypes.most_common(1)[0][0] if weight_dtypes else None
return {'scheme': 'comfy_quant', 'format': COMFY_QUANT_DTYPE_FORMATS.get(dominant), 'marked_layers': len(comfy), 'source': 'weight-dtype'}
has_fp8 = dtypes.get('F8_E4M3', 0) + dtypes.get('F8_E5M2', 0) > 0
has_scales = any(k.endswith(('scaled_fp8', '.scale_weight', '.scale_input', '.weight_scale')) for k in keys)
if has_fp8 and has_scales:
return {'scheme': 'scaled_fp8', 'format': 'float8_e4m3fn' if dtypes.get('F8_E4M3') else 'float8_e5m2', 'marked_layers': None, 'source': 'marker'}
return {'scheme': None, 'format': None, 'marked_layers': None, 'source': None}
def analyze_header(header: dict, container: str = 'safetensors', arch_metadata: dict | None = None) -> dict:
metadata = header.get('__metadata__') or {}
if arch_metadata:
metadata = {**metadata, **arch_metadata}
entries = {k: v for k, v in header.items() if k != '__metadata__' and isinstance(v, dict)}
keys = list(entries)
dtypes = Counter(v.get('dtype') for v in entries.values() if v.get('dtype'))
params = 0
# dominant dtype is element-weighted over core tensors only: bundled
# companions and scale scalars must not outvote the diffusion weights
core_elements = Counter()
for k, v in entries.items():
shape = v.get('shape')
if not isinstance(shape, (list, tuple)) or k.endswith('.comfy_quant'):
continue
n = 1
for dim in shape:
n *= int(dim)
params += n
if not strip_key(k).startswith(COMPANION_PREFIXES) and v.get('dtype'):
core_elements[v['dtype']] += n
inner_keys = [strip_key(k) for k in keys]
shapes = {k: v.get('shape') for k, v in entries.items()}
arch = match_arch(keys, metadata, shapes)
key_dtypes = {strip_key(k): v.get('dtype') for k, v in entries.items()}
quant = detect_quant(inner_keys, dtypes, container, metadata=metadata, key_dtypes=key_dtypes)
flags = []
if not metadata:
flags.append('no_metadata')
if any(k.startswith(COMPANION_PREFIXES) for k in inner_keys):
flags.append('companion_bundled')
if quant['scheme'] == 'comfy_quant':
flags.append('comfy_marker')
if arch['family'] == 'wanai' and arch['variant'] is None:
flags.append('role_ambiguous')
if arch['kind'] == 'lora' and 0 < arch['confidence'] < 0.9:
flags.append('lora_base_inferred')
return {
'schema': SCHEMA_VERSION,
'ok': True,
'error': None,
'container': container,
'tensors': len(keys),
'params': params,
'dtypes': dict(dtypes),
'dominant_dtype': core_elements.most_common(1)[0][0] if core_elements else (dtypes.most_common(1)[0][0] if dtypes else None),
'arch': arch,
'quant': quant,
'metadata': metadata,
'metadata_present': bool(metadata),
'flags': flags,
}
def error_result(container: str, error: str, flag: str = 'unreadable') -> dict:
return {
'schema': SCHEMA_VERSION, 'ok': False, 'error': error, 'container': container,
'tensors': 0, 'params': 0, 'dtypes': {}, 'dominant_dtype': None,
'arch': {'kind': 'unknown', 'family': 'unknown', 'display': 'Unknown', 'confidence': 0.0, 'variant': None, 'detected_prefix': '', 'matched_markers': [], 'candidates': []},
'quant': {'scheme': None, 'format': None, 'marked_layers': None, 'source': None},
'metadata': {}, 'metadata_present': False, 'flags': [flag],
}
def read_safetensors_header(path: str) -> dict:
"""Raw header read: 8-byte length prefix + JSON. Does not touch the global
sd_metadata cache and keeps ss_tag_frequency at full fidelity."""
with open(path, 'rb') as f:
header_len = int.from_bytes(f.read(8), 'little')
if header_len <= 0 or header_len > MAX_HEADER_BYTES:
raise ValueError(f'implausible header length: {header_len}')
return json.loads(f.read(header_len).decode('utf-8'))
def comfy_marker_format(path: str) -> str | None:
"""Exact comfy_quant format string from the first marker tensor's bytes;
only possible for local files, remote peeks stay dtype-inferred."""
try:
with open(path, 'rb') as f:
header_len = int.from_bytes(f.read(8), 'little')
if header_len <= 0 or header_len > MAX_HEADER_BYTES:
return None
header = json.loads(f.read(header_len).decode('utf-8'))
for k, v in header.items():
if k == '__metadata__' or not k.endswith('.comfy_quant') or not isinstance(v, dict):
continue
start, end = v.get('data_offsets', (0, 0))
if end <= start or end - start > 4096:
return None
f.seek(8 + header_len + start)
fmt = json.loads(f.read(end - start).decode('utf-8')).get('format')
return str(fmt) if fmt else None
except Exception:
return None
return None
# quant format / dtype to the precision token used in filenames; fp8 stays
# variant-specific because e4m3fn and e5m2 differ in kernel support
QUANT_PRECISION_TOKENS = {'int8_tensorwise': 'int8', 'float8_e4m3fn': 'fp8_e4m3fn', 'float8_e5m2': 'fp8_e5m2', 'nvfp4': 'nvfp4', 'mxfp8': 'mxfp8'}
DTYPE_PRECISION_TOKENS = {'F32': 'fp32', 'F16': 'fp16', 'BF16': 'bf16', 'F8_E4M3': 'fp8_e4m3fn', 'F8_E5M2': 'fp8_e5m2'}
def precision_token(probe: dict) -> str | None:
"""Filename token for a probe's true precision; None when the container
encodes it already (gguf) or nothing is known."""
quant = probe.get('quant') or {}
scheme = quant.get('scheme')
if scheme == 'gguf':
return None
if scheme in ('comfy_quant', 'scaled_fp8'):
fmt = quant.get('format') or ''
return QUANT_PRECISION_TOKENS.get(fmt, re.sub(r'[^a-z0-9]', '', fmt.lower()) or None)
return DTYPE_PRECISION_TOKENS.get(probe.get('dominant_dtype') or '')
def probe_safetensors_file(path: str) -> dict:
try:
header = read_safetensors_header(path)
except Exception as e:
return error_result('safetensors', str(e), 'corrupt_header')
result = analyze_header(header)
if result['quant']['scheme'] == 'comfy_quant':
fmt = comfy_marker_format(path)
if fmt:
result['quant']['format'] = fmt
result['quant']['source'] = 'marker'
return result
def probe_gguf_file(path: str) -> dict:
try:
from modules.ggml import install_gguf
install_gguf()
import gguf
reader = gguf.GGUFReader(path)
header = {}
for tensor in reader.tensors:
header[str(tensor.name)] = {'dtype': str(tensor.tensor_type.name), 'shape': [int(d) for d in reversed(tensor.shape)]}
arch_metadata = {}
for key in ('general.architecture', 'general.name'):
fld = reader.fields.get(key)
if fld is not None:
try:
arch_metadata[key] = str(bytes(fld.parts[fld.data[0]]).decode('utf-8'))
except Exception:
pass
return analyze_header(header, container='gguf', arch_metadata=arch_metadata)
except Exception as e:
return error_result('gguf', str(e))
probe_cache = None
def probe_file(path: str, use_cache: bool = True) -> dict:
"""Dispatch by extension with an mtime-validated persistent cache."""
global probe_cache # pylint: disable=global-statement
ext = os.path.splitext(path)[1].lower()
if ext not in ('.safetensors', '.gguf'):
return error_result('unknown', f'unsupported extension: {ext}', 'unsupported')
try:
stat = os.stat(path)
except OSError as e:
return error_result('unknown', str(e))
if probe_cache is None:
from modules import paths
from modules.json_helpers import readfile
probe_cache = readfile(paths.probe_cache_file, silent=True, lock=True, as_type='dict')
entry = probe_cache.get(path) if use_cache else None
if entry and entry.get('mtime') == stat.st_mtime and entry.get('size') == stat.st_size and entry.get('schema') == SCHEMA_VERSION:
return entry['probe']
probe = probe_safetensors_file(path) if ext == '.safetensors' else probe_gguf_file(path)
probe_cache[path] = {'mtime': stat.st_mtime, 'size': stat.st_size, 'schema': SCHEMA_VERSION, 'probe': probe}
return probe
def save_probe_cache():
"""Flush the probe cache to disk; call once per scan, not per file."""
if probe_cache is None:
return
from modules import paths
from modules.json_helpers import writefile
writefile(probe_cache, paths.probe_cache_file, silent=True, atomic=True)
modules/paths.py
+2
View File
@@ -33,6 +33,8 @@ data_path = cli.data_dir
models_config = cli.models_dir or config.get('models_dir') or 'models'
models_path = models_config if os.path.isabs(models_config) else os.path.join(data_path, models_config)
params_path = os.environ.get('SD_PATH_PARAMS', os.path.join(data_path, "params.txt"))
probe_cache_file = os.path.join(data_path, "data", "signatures.json")
civitai_probe_file = os.path.join(data_path, "data", "civitai.json")
extensions_dir = cli.extensions_dir or os.path.join(data_path, "extensions")
extensions_builtin_dir = "extensions-builtin"
sd_configs_path = os.path.join(script_path, "configs")
modules/sd_models.py
+11 -4
View File
@@ -331,7 +331,7 @@ def load_diffuser_initial(diffusers_load_config: dict, op='model'):
return sd_model, checkpoint_info
def hf_prefetch_configs(checkpoint_info: CheckpointInfo, diffusers_load_config: dict, op='model'):
def hf_prefetch_configs(checkpoint_info: CheckpointInfo | str, diffusers_load_config: dict, op='model'):
# diffusers pipeline downloads build subfolder config allow-patterns with os.path.join, and huggingface_hub>=1.22
# matches patterns with fnmatchcase which does not normalize separators (huggingface/huggingface_hub#4435),
# so on windows component config.json files are never downloaded and the incomplete snapshot
@@ -496,8 +496,8 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf
sd_model = load_vibe(checkpoint_info, diffusers_load_config)
allow_post_quant = False
elif model_type in ['JoyEdit']:
from pipelines.model_joy import load_joy
sd_model = load_joy(checkpoint_info, diffusers_load_config)
from pipelines.model_joy import load_joyedit
sd_model = load_joyedit(checkpoint_info, diffusers_load_config)
allow_post_quant = False
elif model_type in ['Qwen']:
from pipelines.model_qwen import load_qwen
@@ -932,7 +932,9 @@ def load_diffuser(checkpoint_info: CheckpointInfo | None = None, op='model', rev
if model_type is None:
log.error(f'Load {op}: pipeline={shared.opts.diffusers_pipeline} not detected')
return
hf_prefetch_configs(checkpoint_info, diffusers_load_config, op)
vae_file = None
if model_type.startswith('Stable Diffusion') and (op == 'model' or op == 'refiner'): # preload vae for sd models
vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
@@ -1475,8 +1477,9 @@ def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op='
loaded_ckpt = getattr(sd_model, 'sd_checkpoint_info', None) if sd_model is not None else None
changed_checkpoint = loaded_ckpt is None or checkpoint_info is None or loaded_ckpt.filename != checkpoint_info.filename
reset_unet = shared.opts.sd_unet not in (None, 'Default', 'None')
reset_unet_secondary = shared.opts.sd_unet_secondary not in (None, 'Default', 'None')
reset_te = shared.opts.sd_text_encoder not in (None, 'Default', 'None')
if op == 'model' and sd_model is not None and changed_checkpoint and (reset_unet or reset_te):
if op == 'model' and sd_model is not None and changed_checkpoint and (reset_unet or reset_unet_secondary or reset_te):
# compare detected model type, not pipeline class: custom-loader arches (e.g. Krea2) load as a
# concrete class but detect as generic DiffusionPipeline, so a class compare would falsely reset
# across same-arch checkpoints (Base vs Turbo). detect both sides so the comparison is symmetric.
@@ -1490,6 +1493,10 @@ def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op='
log.info(f'Load model: type="{old_type}" changed="{new_type}" unet="{shared.opts.sd_unet}" set to default')
shared.opts.data["sd_unet"] = 'Default'
sd_unet.loaded_unet = None
if reset_unet_secondary:
log.info(f'Load model: type="{old_type}" changed="{new_type}" unet_secondary="{shared.opts.sd_unet_secondary}" set to default')
shared.opts.data["sd_unet_secondary"] = 'Default'
sd_unet.loaded_unet_secondary = None
if reset_te:
log.info(f'Load model: type="{old_type}" changed="{new_type}" te="{shared.opts.sd_text_encoder}" set to default')
shared.opts.data["sd_text_encoder"] = 'Default'
modules/sd_models_utils.py
+1 -1
View File
@@ -58,7 +58,7 @@ def get_call(cls):
return signature.parameters
def path_to_repo(checkpoint_info):
def path_to_repo(checkpoint_info: CheckpointInfo | str):
if isinstance(checkpoint_info, CheckpointInfo):
if os.path.exists(checkpoint_info.path) and 'models--' not in checkpoint_info.path:
return checkpoint_info.path # local models
modules/sd_offload.py
+2 -2
View File
@@ -21,8 +21,8 @@ offload_post = ['h1']
offload_hook_instance = None
balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline']
no_split_module_classes = [
"Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d",
"SDNQLinear", "SDNQConv1d", "SDNQConv2d", "SDNQConv3d", "SDNQConvTranspose1d", "SDNQConvTranspose2d", "SDNQConvTranspose3d",
"Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "Embedding",
"SDNQLinear", "SDNQConv1d", "SDNQConv2d", "SDNQConv3d", "SDNQConvTranspose1d", "SDNQConvTranspose2d", "SDNQConvTranspose3d", "SDNQEmbedding",
"WanTransformerBlock",
]
accelerate_dtype_byte_size = None
modules/sd_unet.py
+32
View File
@@ -5,11 +5,14 @@ from modules.logger import log
unet_dict = {}
loaded_unet = None
loaded_unet_secondary = None
failed_unet = []
debug = os.environ.get('SD_LOAD_DEBUG', None) is not None
dit_models = ['Flux', 'StableDiffusion3', 'HiDream', 'Lumina2', 'Chroma', 'Wan', 'Qwen', 'Anima']
# model types (shared.sd_model_type keyspace) the secondary UNET override applies to
DUAL_TRANSFORMER_TYPES = ('ideogram4', 'wanai')
def load_unet_sdxl_nunchaku(repo_id):
@@ -101,6 +104,35 @@ def load_unet(model, repo_id: str | None = None):
devices.torch_gc()
def load_unet_secondary(model): # pylint: disable=unused-argument
"""Onchange handler for the secondary UNET override: a change means a
full reload; for single-transformer models the selection is stored and
applies on the next dual-transformer load.
"""
global loaded_unet_secondary # pylint: disable=global-statement
selected = shared.opts.sd_unet_secondary
if selected is None or selected in ('Default', 'None'):
if loaded_unet_secondary in (None, 'Default', 'None'):
return
log.info(f'Load module: type=UNet slot=secondary name="Default" (was="{loaded_unet_secondary}") reverting to base transformer')
loaded_unet_secondary = selected
sd_models.reload_model_weights(force=True)
return
if selected not in list(unet_dict):
log.error(f'Load module: type=UNet slot=secondary not found: {selected}')
return
if selected == loaded_unet_secondary or selected in failed_unet:
return
if shared.sd_model_type not in DUAL_TRANSFORMER_TYPES:
log.warning(f'Load module: type=UNet slot=secondary name="{selected}" stored: model type={shared.sd_model_type} has a single transformer, applies on next dual-transformer load')
return
loaded_unet_secondary = selected
sd_models.reload_model_weights(force=True)
devices.torch_gc()
def refresh_unet_list():
unet_dict.clear()
for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors", ".gguf", ".pth"]):
modules/sdnq/common.py
+137 -223
View File
@@ -4,7 +4,7 @@ import os
import json
import torch
from modules import shared, devices
from modules import shared
sdnq_version = "0.2.2"
sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"}
@@ -62,171 +62,171 @@ dtype_dict = {
"float8_e4m3fn": {"min": -448.0, "max": 448.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False, "is_packed": False},
"float8_e5m2": {"min": -57344.0, "max": 57344.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False, "is_packed": False},
### Custom Floats
"float16_e1m14fn": {"min": -3.9998779296875, "max": 3.9998779296875, "num_bits": 16, "sign": 1, "exponent": 1, "mantissa": 14, "min_normal": 1.00006103515625, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e2m13fn": {"min": -7.99951171875, "max": 7.99951171875, "num_bits": 16, "sign": 1, "exponent": 2, "mantissa": 13, "min_normal": 0.50006103515625, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e3m12fn": {"min": -31.99609375, "max": 31.99609375, "num_bits": 16, "sign": 1, "exponent": 3, "mantissa": 12, "min_normal": 0.125030517578125, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e4m11fn": {"min": -511.875, "max": 511.875, "num_bits": 16, "sign": 1, "exponent": 4, "mantissa": 11, "min_normal": 0.007816314697265625, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e5m10fn": {"min": -131008.0, "max": 131008.0, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "min_normal": 3.0547380447387695e-05, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e1m14fn": {"min": -3.9998779296875, "max": 3.9998779296875, "num_bits": 16, "sign": 1, "exponent": 1, "mantissa": 14, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e2m13fn": {"min": -7.99951171875, "max": 7.99951171875, "num_bits": 16, "sign": 1, "exponent": 2, "mantissa": 13, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e3m12fn": {"min": -31.99609375, "max": 31.99609375, "num_bits": 16, "sign": 1, "exponent": 3, "mantissa": 12, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e4m11fn": {"min": -511.875, "max": 511.875, "num_bits": 16, "sign": 1, "exponent": 4, "mantissa": 11, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e5m10fn": {"min": -131008.0, "max": 131008.0, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float15_e1m13fn": {"min": -3.999755859375, "max": 3.999755859375, "num_bits": 15, "sign": 1, "exponent": 1, "mantissa": 13, "min_normal": 1.0001220703125, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e2m12fn": {"min": -7.9990234375, "max": 7.9990234375, "num_bits": 15, "sign": 1, "exponent": 2, "mantissa": 12, "min_normal": 0.5001220703125, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e3m11fn": {"min": -31.9921875, "max": 31.9921875, "num_bits": 15, "sign": 1, "exponent": 3, "mantissa": 11, "min_normal": 0.12506103515625, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e4m10fn": {"min": -511.75, "max": 511.75, "num_bits": 15, "sign": 1, "exponent": 4, "mantissa": 10, "min_normal": 0.00782012939453125, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e5m9fn": {"min": -130944.0, "max": 130944.0, "num_bits": 15, "sign": 1, "exponent": 5, "mantissa": 9, "min_normal": 3.057718276977539e-05, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e1m13fn": {"min": -3.999755859375, "max": 3.999755859375, "num_bits": 15, "sign": 1, "exponent": 1, "mantissa": 13, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e2m12fn": {"min": -7.9990234375, "max": 7.9990234375, "num_bits": 15, "sign": 1, "exponent": 2, "mantissa": 12, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e3m11fn": {"min": -31.9921875, "max": 31.9921875, "num_bits": 15, "sign": 1, "exponent": 3, "mantissa": 11, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e4m10fn": {"min": -511.75, "max": 511.75, "num_bits": 15, "sign": 1, "exponent": 4, "mantissa": 10, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e5m9fn": {"min": -130944.0, "max": 130944.0, "num_bits": 15, "sign": 1, "exponent": 5, "mantissa": 9, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float14_e1m12fn": {"min": -3.99951171875, "max": 3.99951171875, "num_bits": 14, "sign": 1, "exponent": 1, "mantissa": 12, "min_normal": 1.000244140625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e2m11fn": {"min": -7.998046875, "max": 7.998046875, "num_bits": 14, "sign": 1, "exponent": 2, "mantissa": 11, "min_normal": 0.500244140625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e3m10fn": {"min": -31.984375, "max": 31.984375, "num_bits": 14, "sign": 1, "exponent": 3, "mantissa": 10, "min_normal": 0.1251220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e4m9fn": {"min": -511.5, "max": 511.5, "num_bits": 14, "sign": 1, "exponent": 4, "mantissa": 9, "min_normal": 0.0078277587890625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e5m8fn": {"min": -130816.0, "max": 130816.0, "num_bits": 14, "sign": 1, "exponent": 5, "mantissa": 8, "min_normal": 3.063678741455078e-05, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e1m12fn": {"min": -3.99951171875, "max": 3.99951171875, "num_bits": 14, "sign": 1, "exponent": 1, "mantissa": 12, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e2m11fn": {"min": -7.998046875, "max": 7.998046875, "num_bits": 14, "sign": 1, "exponent": 2, "mantissa": 11, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e3m10fn": {"min": -31.984375, "max": 31.984375, "num_bits": 14, "sign": 1, "exponent": 3, "mantissa": 10, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e4m9fn": {"min": -511.5, "max": 511.5, "num_bits": 14, "sign": 1, "exponent": 4, "mantissa": 9, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e5m8fn": {"min": -130816.0, "max": 130816.0, "num_bits": 14, "sign": 1, "exponent": 5, "mantissa": 8, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float13_e1m11fn": {"min": -3.9990234375, "max": 3.9990234375, "num_bits": 13, "sign": 1, "exponent": 1, "mantissa": 11, "min_normal": 1.00048828125, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e2m10fn": {"min": -7.99609375, "max": 7.99609375, "num_bits": 13, "sign": 1, "exponent": 2, "mantissa": 10, "min_normal": 0.50048828125, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e3m9fn": {"min": -31.96875, "max": 31.96875, "num_bits": 13, "sign": 1, "exponent": 3, "mantissa": 9, "min_normal": 0.125244140625, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e4m8fn": {"min": -511.0, "max": 511.0, "num_bits": 13, "sign": 1, "exponent": 4, "mantissa": 8, "min_normal": 0.007843017578125, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e5m7fn": {"min": -130560.0, "max": 130560.0, "num_bits": 13, "sign": 1, "exponent": 5, "mantissa": 7, "min_normal": 3.075599670410156e-05, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e1m11fn": {"min": -3.9990234375, "max": 3.9990234375, "num_bits": 13, "sign": 1, "exponent": 1, "mantissa": 11, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e2m10fn": {"min": -7.99609375, "max": 7.99609375, "num_bits": 13, "sign": 1, "exponent": 2, "mantissa": 10, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e3m9fn": {"min": -31.96875, "max": 31.96875, "num_bits": 13, "sign": 1, "exponent": 3, "mantissa": 9, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e4m8fn": {"min": -511.0, "max": 511.0, "num_bits": 13, "sign": 1, "exponent": 4, "mantissa": 8, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e5m7fn": {"min": -130560.0, "max": 130560.0, "num_bits": 13, "sign": 1, "exponent": 5, "mantissa": 7, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float12_e1m10fn": {"min": -3.998046875, "max": 3.998046875, "num_bits": 12, "sign": 1, "exponent": 1, "mantissa": 10, "min_normal": 1.0009765625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e2m9fn": {"min": -7.9921875, "max": 7.9921875, "num_bits": 12, "sign": 1, "exponent": 2, "mantissa": 9, "min_normal": 0.5009765625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e3m8fn": {"min": -31.9375, "max": 31.9375, "num_bits": 12, "sign": 1, "exponent": 3, "mantissa": 8, "min_normal": 0.12548828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e4m7fn": {"min": -510.0, "max": 510.0, "num_bits": 12, "sign": 1, "exponent": 4, "mantissa": 7, "min_normal": 0.00787353515625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e5m6fn": {"min": -130048.0, "max": 130048.0, "num_bits": 12, "sign": 1, "exponent": 5, "mantissa": 6, "min_normal": 3.0994415283203125e-05, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e1m10fn": {"min": -3.998046875, "max": 3.998046875, "num_bits": 12, "sign": 1, "exponent": 1, "mantissa": 10, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e2m9fn": {"min": -7.9921875, "max": 7.9921875, "num_bits": 12, "sign": 1, "exponent": 2, "mantissa": 9, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e3m8fn": {"min": -31.9375, "max": 31.9375, "num_bits": 12, "sign": 1, "exponent": 3, "mantissa": 8, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e4m7fn": {"min": -510.0, "max": 510.0, "num_bits": 12, "sign": 1, "exponent": 4, "mantissa": 7, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e5m6fn": {"min": -130048.0, "max": 130048.0, "num_bits": 12, "sign": 1, "exponent": 5, "mantissa": 6, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float11_e1m9fn": {"min": -3.99609375, "max": 3.99609375, "num_bits": 11, "sign": 1, "exponent": 1, "mantissa": 9, "min_normal": 1.001953125, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e2m8fn": {"min": -7.984375, "max": 7.984375, "num_bits": 11, "sign": 1, "exponent": 2, "mantissa": 8, "min_normal": 0.501953125, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e3m7fn": {"min": -31.875, "max": 31.875, "num_bits": 11, "sign": 1, "exponent": 3, "mantissa": 7, "min_normal": 0.1259765625, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e4m6fn": {"min": -508.0, "max": 508.0, "num_bits": 11, "sign": 1, "exponent": 4, "mantissa": 6, "min_normal": 0.0079345703125, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e5m5fn": {"min": -129024.0, "max": 129024.0, "num_bits": 11, "sign": 1, "exponent": 5, "mantissa": 5, "min_normal": 3.147125244140625e-05, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e1m9fn": {"min": -3.99609375, "max": 3.99609375, "num_bits": 11, "sign": 1, "exponent": 1, "mantissa": 9, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e2m8fn": {"min": -7.984375, "max": 7.984375, "num_bits": 11, "sign": 1, "exponent": 2, "mantissa": 8, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e3m7fn": {"min": -31.875, "max": 31.875, "num_bits": 11, "sign": 1, "exponent": 3, "mantissa": 7, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e4m6fn": {"min": -508.0, "max": 508.0, "num_bits": 11, "sign": 1, "exponent": 4, "mantissa": 6, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e5m5fn": {"min": -129024.0, "max": 129024.0, "num_bits": 11, "sign": 1, "exponent": 5, "mantissa": 5, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float10_e1m8fn": {"min": -3.9921875, "max": 3.9921875, "num_bits": 10, "sign": 1, "exponent": 1, "mantissa": 8, "min_normal": 1.00390625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e2m7fn": {"min": -7.96875, "max": 7.96875, "num_bits": 10, "sign": 1, "exponent": 2, "mantissa": 7, "min_normal": 0.50390625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e3m6fn": {"min": -31.75, "max": 31.75, "num_bits": 10, "sign": 1, "exponent": 3, "mantissa": 6, "min_normal": 0.126953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e4m5fn": {"min": -504.0, "max": 504.0, "num_bits": 10, "sign": 1, "exponent": 4, "mantissa": 5, "min_normal": 0.008056640625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e5m4fn": {"min": -126976.0, "max": 126976.0, "num_bits": 10, "sign": 1, "exponent": 5, "mantissa": 4, "min_normal": 3.24249267578125e-05, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e1m8fn": {"min": -3.9921875, "max": 3.9921875, "num_bits": 10, "sign": 1, "exponent": 1, "mantissa": 8, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e2m7fn": {"min": -7.96875, "max": 7.96875, "num_bits": 10, "sign": 1, "exponent": 2, "mantissa": 7, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e3m6fn": {"min": -31.75, "max": 31.75, "num_bits": 10, "sign": 1, "exponent": 3, "mantissa": 6, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e4m5fn": {"min": -504.0, "max": 504.0, "num_bits": 10, "sign": 1, "exponent": 4, "mantissa": 5, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e5m4fn": {"min": -126976.0, "max": 126976.0, "num_bits": 10, "sign": 1, "exponent": 5, "mantissa": 4, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float9_e1m7fn": {"min": -3.984375, "max": 3.984375, "num_bits": 9, "sign": 1, "exponent": 1, "mantissa": 7, "min_normal": 1.0078125, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e2m6fn": {"min": -7.9375, "max": 7.9375, "num_bits": 9, "sign": 1, "exponent": 2, "mantissa": 6, "min_normal": 0.5078125, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e3m5fn": {"min": -31.5, "max": 31.5, "num_bits": 9, "sign": 1, "exponent": 3, "mantissa": 5, "min_normal": 0.12890625, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e4m4fn": {"min": -496.0, "max": 496.0, "num_bits": 9, "sign": 1, "exponent": 4, "mantissa": 4, "min_normal": 0.00830078125, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e5m3fn": {"min": -122880.0, "max": 122880.0, "num_bits": 9, "sign": 1, "exponent": 5, "mantissa": 3, "min_normal": 3.4332275390625e-05, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e1m7fn": {"min": -3.984375, "max": 3.984375, "num_bits": 9, "sign": 1, "exponent": 1, "mantissa": 7, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e2m6fn": {"min": -7.9375, "max": 7.9375, "num_bits": 9, "sign": 1, "exponent": 2, "mantissa": 6, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e3m5fn": {"min": -31.5, "max": 31.5, "num_bits": 9, "sign": 1, "exponent": 3, "mantissa": 5, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e4m4fn": {"min": -496.0, "max": 496.0, "num_bits": 9, "sign": 1, "exponent": 4, "mantissa": 4, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e5m3fn": {"min": -122880.0, "max": 122880.0, "num_bits": 9, "sign": 1, "exponent": 5, "mantissa": 3, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float8_e1m6fn": {"min": -3.96875, "max": 3.96875, "num_bits": 8, "sign": 1, "exponent": 1, "mantissa": 6, "min_normal": 1.015625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e2m5fn": {"min": -7.875, "max": 7.875, "num_bits": 8, "sign": 1, "exponent": 2, "mantissa": 5, "min_normal": 0.515625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e3m4fn": {"min": -31.0, "max": 31.0, "num_bits": 8, "sign": 1, "exponent": 3, "mantissa": 4, "min_normal": 0.1328125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e4m3fn_sdnq": {"min": -480.0, "max": 480.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "min_normal": 0.0087890625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e5m2fn": {"min": -114688.0, "max": 114688.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "min_normal": 3.814697265625e-05, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e1m6fn": {"min": -3.96875, "max": 3.96875, "num_bits": 8, "sign": 1, "exponent": 1, "mantissa": 6, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e2m5fn": {"min": -7.875, "max": 7.875, "num_bits": 8, "sign": 1, "exponent": 2, "mantissa": 5, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e3m4fn": {"min": -31.0, "max": 31.0, "num_bits": 8, "sign": 1, "exponent": 3, "mantissa": 4, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e4m3fn_sdnq": {"min": -480.0, "max": 480.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e5m2fn": {"min": -114688.0, "max": 114688.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float7_e1m5fn": {"min": -3.9375, "max": 3.9375, "num_bits": 7, "sign": 1, "exponent": 1, "mantissa": 5, "min_normal": 1.03125, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e2m4fn": {"min": -7.75, "max": 7.75, "num_bits": 7, "sign": 1, "exponent": 2, "mantissa": 4, "min_normal": 0.53125, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e3m3fn": {"min": -30.0, "max": 30.0, "num_bits": 7, "sign": 1, "exponent": 3, "mantissa": 3, "min_normal": 0.140625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e4m2fn": {"min": -448.0, "max": 448.0, "num_bits": 7, "sign": 1, "exponent": 4, "mantissa": 2, "min_normal": 0.009765625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e5m1fn": {"min": -98304.0, "max": 98304.0, "num_bits": 7, "sign": 1, "exponent": 5, "mantissa": 1, "min_normal": 4.57763671875e-05, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e1m5fn": {"min": -3.9375, "max": 3.9375, "num_bits": 7, "sign": 1, "exponent": 1, "mantissa": 5, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e2m4fn": {"min": -7.75, "max": 7.75, "num_bits": 7, "sign": 1, "exponent": 2, "mantissa": 4, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e3m3fn": {"min": -30.0, "max": 30.0, "num_bits": 7, "sign": 1, "exponent": 3, "mantissa": 3, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e4m2fn": {"min": -448.0, "max": 448.0, "num_bits": 7, "sign": 1, "exponent": 4, "mantissa": 2, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e5m1fn": {"min": -98304.0, "max": 98304.0, "num_bits": 7, "sign": 1, "exponent": 5, "mantissa": 1, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float6_e1m4fn": {"min": -3.875, "max": 3.875, "num_bits": 6, "sign": 1, "exponent": 1, "mantissa": 4, "min_normal": 1.0625, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e2m3fn": {"min": -7.5, "max": 7.5, "num_bits": 6, "sign": 1, "exponent": 2, "mantissa": 3, "min_normal": 0.5625, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e3m2fn": {"min": -28.0, "max": 28.0, "num_bits": 6, "sign": 1, "exponent": 3, "mantissa": 2, "min_normal": 0.15625, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e4m1fn": {"min": -384.0, "max": 384.0, "num_bits": 6, "sign": 1, "exponent": 4, "mantissa": 1, "min_normal": 0.01171875, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e5m0fn": {"min": -65536.0, "max": 65536.0, "num_bits": 6, "sign": 1, "exponent": 5, "mantissa": 0, "min_normal": 6.103515625e-05, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e1m4fn": {"min": -3.875, "max": 3.875, "num_bits": 6, "sign": 1, "exponent": 1, "mantissa": 4, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e2m3fn": {"min": -7.5, "max": 7.5, "num_bits": 6, "sign": 1, "exponent": 2, "mantissa": 3, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e3m2fn": {"min": -28.0, "max": 28.0, "num_bits": 6, "sign": 1, "exponent": 3, "mantissa": 2, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e4m1fn": {"min": -384.0, "max": 384.0, "num_bits": 6, "sign": 1, "exponent": 4, "mantissa": 1, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e5m0fn": {"min": -65536.0, "max": 65536.0, "num_bits": 6, "sign": 1, "exponent": 5, "mantissa": 0, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float5_e1m3fn": {"min": -3.75, "max": 3.75, "num_bits": 5, "sign": 1, "exponent": 1, "mantissa": 3, "min_normal": 1.125, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e2m2fn": {"min": -7.0, "max": 7.0, "num_bits": 5, "sign": 1, "exponent": 2, "mantissa": 2, "min_normal": 0.625, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e3m1fn": {"min": -24.0, "max": 24.0, "num_bits": 5, "sign": 1, "exponent": 3, "mantissa": 1, "min_normal": 0.1875, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e4m0fn": {"min": -256.0, "max": 256.0, "num_bits": 5, "sign": 1, "exponent": 4, "mantissa": 0, "min_normal": 0.015625, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e1m3fn": {"min": -3.75, "max": 3.75, "num_bits": 5, "sign": 1, "exponent": 1, "mantissa": 3, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e2m2fn": {"min": -7.0, "max": 7.0, "num_bits": 5, "sign": 1, "exponent": 2, "mantissa": 2, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e3m1fn": {"min": -24.0, "max": 24.0, "num_bits": 5, "sign": 1, "exponent": 3, "mantissa": 1, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e4m0fn": {"min": -256.0, "max": 256.0, "num_bits": 5, "sign": 1, "exponent": 4, "mantissa": 0, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float4_e1m2fn": {"min": -3.5, "max": 3.5, "num_bits": 4, "sign": 1, "exponent": 1, "mantissa": 2, "min_normal": 1.25, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e2m1fn": {"min": -6.0, "max": 6.0, "num_bits": 4, "sign": 1, "exponent": 2, "mantissa": 1, "min_normal": 0.75, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e3m0fn": {"min": -16.0, "max": 16.0, "num_bits": 4, "sign": 1, "exponent": 3, "mantissa": 0, "min_normal": 0.25, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e1m2fn": {"min": -3.5, "max": 3.5, "num_bits": 4, "sign": 1, "exponent": 1, "mantissa": 2, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e2m1fn": {"min": -6.0, "max": 6.0, "num_bits": 4, "sign": 1, "exponent": 2, "mantissa": 1, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e3m0fn": {"min": -16.0, "max": 16.0, "num_bits": 4, "sign": 1, "exponent": 3, "mantissa": 0, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float3_e1m1fn": {"min": -3.0, "max": 3.0, "num_bits": 3, "sign": 1, "exponent": 1, "mantissa": 1, "min_normal": 1.5, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float3_e2m0fn": {"min": -4.0, "max": 4.0, "num_bits": 3, "sign": 1, "exponent": 2, "mantissa": 0, "min_normal": 1.0, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float3_e1m1fn": {"min": -3.0, "max": 3.0, "num_bits": 3, "sign": 1, "exponent": 1, "mantissa": 1, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float3_e2m0fn": {"min": -4.0, "max": 4.0, "num_bits": 3, "sign": 1, "exponent": 2, "mantissa": 0, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float2_e1m0fn": {"min": -2.0, "max": 2.0, "num_bits": 2, "sign": 1, "exponent": 1, "mantissa": 0, "min_normal": 2.0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float2_e1m0fn": {"min": -2.0, "max": 2.0, "num_bits": 2, "sign": 1, "exponent": 1, "mantissa": 0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
### Custom Unsigned Floats
"float16_e1m15fnu": {"min": 0, "max": 3.99993896484375, "num_bits": 16, "sign": 0, "exponent": 1, "mantissa": 15, "min_normal": 1.000030517578125, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e2m14fnu": {"min": 0, "max": 7.999755859375, "num_bits": 16, "sign": 0, "exponent": 2, "mantissa": 14, "min_normal": 0.500030517578125, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e3m13fnu": {"min": 0, "max": 31.998046875, "num_bits": 16, "sign": 0, "exponent": 3, "mantissa": 13, "min_normal": 0.1250152587890625, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e4m12fnu": {"min": 0, "max": 511.9375, "num_bits": 16, "sign": 0, "exponent": 4, "mantissa": 12, "min_normal": 0.007814407348632812, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e5m11fnu": {"min": 0, "max": 131040.0, "num_bits": 16, "sign": 0, "exponent": 5, "mantissa": 11, "min_normal": 3.053247928619385e-05, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e1m15fnu": {"min": 0, "max": 3.99993896484375, "num_bits": 16, "sign": 0, "exponent": 1, "mantissa": 15, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e2m14fnu": {"min": 0, "max": 7.999755859375, "num_bits": 16, "sign": 0, "exponent": 2, "mantissa": 14, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e3m13fnu": {"min": 0, "max": 31.998046875, "num_bits": 16, "sign": 0, "exponent": 3, "mantissa": 13, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e4m12fnu": {"min": 0, "max": 511.9375, "num_bits": 16, "sign": 0, "exponent": 4, "mantissa": 12, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e5m11fnu": {"min": 0, "max": 131040.0, "num_bits": 16, "sign": 0, "exponent": 5, "mantissa": 11, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float15_e1m14fnu": {"min": 0, "max": 3.9998779296875, "num_bits": 15, "sign": 0, "exponent": 1, "mantissa": 14, "min_normal": 1.00006103515625, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e2m13fnu": {"min": 0, "max": 7.99951171875, "num_bits": 15, "sign": 0, "exponent": 2, "mantissa": 13, "min_normal": 0.50006103515625, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e3m12fnu": {"min": 0, "max": 31.99609375, "num_bits": 15, "sign": 0, "exponent": 3, "mantissa": 12, "min_normal": 0.125030517578125, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e4m11fnu": {"min": 0, "max": 511.875, "num_bits": 15, "sign": 0, "exponent": 4, "mantissa": 11, "min_normal": 0.007816314697265625, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e5m10fnu": {"min": 0, "max": 131008.0, "num_bits": 15, "sign": 0, "exponent": 5, "mantissa": 10, "min_normal": 3.0547380447387695e-05, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e1m14fnu": {"min": 0, "max": 3.9998779296875, "num_bits": 15, "sign": 0, "exponent": 1, "mantissa": 14, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e2m13fnu": {"min": 0, "max": 7.99951171875, "num_bits": 15, "sign": 0, "exponent": 2, "mantissa": 13, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e3m12fnu": {"min": 0, "max": 31.99609375, "num_bits": 15, "sign": 0, "exponent": 3, "mantissa": 12, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e4m11fnu": {"min": 0, "max": 511.875, "num_bits": 15, "sign": 0, "exponent": 4, "mantissa": 11, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e5m10fnu": {"min": 0, "max": 131008.0, "num_bits": 15, "sign": 0, "exponent": 5, "mantissa": 10, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float14_e1m13fnu": {"min": 0, "max": 3.999755859375, "num_bits": 14, "sign": 0, "exponent": 1, "mantissa": 13, "min_normal": 1.0001220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e2m12fnu": {"min": 0, "max": 7.9990234375, "num_bits": 14, "sign": 0, "exponent": 2, "mantissa": 12, "min_normal": 0.5001220703125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e3m11fnu": {"min": 0, "max": 31.9921875, "num_bits": 14, "sign": 0, "exponent": 3, "mantissa": 11, "min_normal": 0.12506103515625, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e4m10fnu": {"min": 0, "max": 511.75, "num_bits": 14, "sign": 0, "exponent": 4, "mantissa": 10, "min_normal": 0.00782012939453125, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e5m9fnu": {"min": 0, "max": 130944.0, "num_bits": 14, "sign": 0, "exponent": 5, "mantissa": 9, "min_normal": 3.057718276977539e-05, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e1m13fnu": {"min": 0, "max": 3.999755859375, "num_bits": 14, "sign": 0, "exponent": 1, "mantissa": 13, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e2m12fnu": {"min": 0, "max": 7.9990234375, "num_bits": 14, "sign": 0, "exponent": 2, "mantissa": 12, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e3m11fnu": {"min": 0, "max": 31.9921875, "num_bits": 14, "sign": 0, "exponent": 3, "mantissa": 11, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e4m10fnu": {"min": 0, "max": 511.75, "num_bits": 14, "sign": 0, "exponent": 4, "mantissa": 10, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e5m9fnu": {"min": 0, "max": 130944.0, "num_bits": 14, "sign": 0, "exponent": 5, "mantissa": 9, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float13_e1m12fnu": {"min": 0, "max": 3.99951171875, "num_bits": 13, "sign": 0, "exponent": 1, "mantissa": 12, "min_normal": 1.000244140625, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e2m11fnu": {"min": 0, "max": 7.998046875, "num_bits": 13, "sign": 0, "exponent": 2, "mantissa": 11, "min_normal": 0.500244140625, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e3m10fnu": {"min": 0, "max": 31.984375, "num_bits": 13, "sign": 0, "exponent": 3, "mantissa": 10, "min_normal": 0.1251220703125, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e4m9fnu": {"min": 0, "max": 511.5, "num_bits": 13, "sign": 0, "exponent": 4, "mantissa": 9, "min_normal": 0.0078277587890625, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e5m8fnu": {"min": 0, "max": 130816.0, "num_bits": 13, "sign": 0, "exponent": 5, "mantissa": 8, "min_normal": 3.063678741455078e-05, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e1m12fnu": {"min": 0, "max": 3.99951171875, "num_bits": 13, "sign": 0, "exponent": 1, "mantissa": 12, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e2m11fnu": {"min": 0, "max": 7.998046875, "num_bits": 13, "sign": 0, "exponent": 2, "mantissa": 11, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e3m10fnu": {"min": 0, "max": 31.984375, "num_bits": 13, "sign": 0, "exponent": 3, "mantissa": 10, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e4m9fnu": {"min": 0, "max": 511.5, "num_bits": 13, "sign": 0, "exponent": 4, "mantissa": 9, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e5m8fnu": {"min": 0, "max": 130816.0, "num_bits": 13, "sign": 0, "exponent": 5, "mantissa": 8, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float12_e1m11fnu": {"min": 0, "max": 3.9990234375, "num_bits": 12, "sign": 0, "exponent": 1, "mantissa": 11, "min_normal": 1.00048828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e2m10fnu": {"min": 0, "max": 7.99609375, "num_bits": 12, "sign": 0, "exponent": 2, "mantissa": 10, "min_normal": 0.50048828125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e3m9fnu": {"min": 0, "max": 31.96875, "num_bits": 12, "sign": 0, "exponent": 3, "mantissa": 9, "min_normal": 0.125244140625, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e4m8fnu": {"min": 0, "max": 511.0, "num_bits": 12, "sign": 0, "exponent": 4, "mantissa": 8, "min_normal": 0.007843017578125, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e5m7fnu": {"min": 0, "max": 130560.0, "num_bits": 12, "sign": 0, "exponent": 5, "mantissa": 7, "min_normal": 3.075599670410156e-05, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e1m11fnu": {"min": 0, "max": 3.9990234375, "num_bits": 12, "sign": 0, "exponent": 1, "mantissa": 11, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e2m10fnu": {"min": 0, "max": 7.99609375, "num_bits": 12, "sign": 0, "exponent": 2, "mantissa": 10, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e3m9fnu": {"min": 0, "max": 31.96875, "num_bits": 12, "sign": 0, "exponent": 3, "mantissa": 9, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e4m8fnu": {"min": 0, "max": 511.0, "num_bits": 12, "sign": 0, "exponent": 4, "mantissa": 8, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e5m7fnu": {"min": 0, "max": 130560.0, "num_bits": 12, "sign": 0, "exponent": 5, "mantissa": 7, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float11_e1m10fnu": {"min": 0, "max": 3.998046875, "num_bits": 11, "sign": 0, "exponent": 1, "mantissa": 10, "min_normal": 1.0009765625, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e2m9fnu": {"min": 0, "max": 7.9921875, "num_bits": 11, "sign": 0, "exponent": 2, "mantissa": 9, "min_normal": 0.5009765625, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e3m8fnu": {"min": 0, "max": 31.9375, "num_bits": 11, "sign": 0, "exponent": 3, "mantissa": 8, "min_normal": 0.12548828125, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e4m7fnu": {"min": 0, "max": 510.0, "num_bits": 11, "sign": 0, "exponent": 4, "mantissa": 7, "min_normal": 0.00787353515625, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e5m6fnu": {"min": 0, "max": 130048.0, "num_bits": 11, "sign": 0, "exponent": 5, "mantissa": 6, "min_normal": 3.0994415283203125e-05, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e1m10fnu": {"min": 0, "max": 3.998046875, "num_bits": 11, "sign": 0, "exponent": 1, "mantissa": 10, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e2m9fnu": {"min": 0, "max": 7.9921875, "num_bits": 11, "sign": 0, "exponent": 2, "mantissa": 9, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e3m8fnu": {"min": 0, "max": 31.9375, "num_bits": 11, "sign": 0, "exponent": 3, "mantissa": 8, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e4m7fnu": {"min": 0, "max": 510.0, "num_bits": 11, "sign": 0, "exponent": 4, "mantissa": 7, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e5m6fnu": {"min": 0, "max": 130048.0, "num_bits": 11, "sign": 0, "exponent": 5, "mantissa": 6, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float10_e1m9fnu": {"min": 0, "max": 3.99609375, "num_bits": 10, "sign": 0, "exponent": 1, "mantissa": 9, "min_normal": 1.001953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e2m8fnu": {"min": 0, "max": 7.984375, "num_bits": 10, "sign": 0, "exponent": 2, "mantissa": 8, "min_normal": 0.501953125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e3m7fnu": {"min": 0, "max": 31.875, "num_bits": 10, "sign": 0, "exponent": 3, "mantissa": 7, "min_normal": 0.1259765625, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e4m6fnu": {"min": 0, "max": 508.0, "num_bits": 10, "sign": 0, "exponent": 4, "mantissa": 6, "min_normal": 0.0079345703125, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e5m5fnu": {"min": 0, "max": 129024.0, "num_bits": 10, "sign": 0, "exponent": 5, "mantissa": 5, "min_normal": 3.147125244140625e-05, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e1m9fnu": {"min": 0, "max": 3.99609375, "num_bits": 10, "sign": 0, "exponent": 1, "mantissa": 9, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e2m8fnu": {"min": 0, "max": 7.984375, "num_bits": 10, "sign": 0, "exponent": 2, "mantissa": 8, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e3m7fnu": {"min": 0, "max": 31.875, "num_bits": 10, "sign": 0, "exponent": 3, "mantissa": 7, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e4m6fnu": {"min": 0, "max": 508.0, "num_bits": 10, "sign": 0, "exponent": 4, "mantissa": 6, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e5m5fnu": {"min": 0, "max": 129024.0, "num_bits": 10, "sign": 0, "exponent": 5, "mantissa": 5, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float9_e1m8fnu": {"min": 0, "max": 3.9921875, "num_bits": 9, "sign": 0, "exponent": 1, "mantissa": 8, "min_normal": 1.00390625, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e2m7fnu": {"min": 0, "max": 7.96875, "num_bits": 9, "sign": 0, "exponent": 2, "mantissa": 7, "min_normal": 0.50390625, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e3m6fnu": {"min": 0, "max": 31.75, "num_bits": 9, "sign": 0, "exponent": 3, "mantissa": 6, "min_normal": 0.126953125, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e4m5fnu": {"min": 0, "max": 504.0, "num_bits": 9, "sign": 0, "exponent": 4, "mantissa": 5, "min_normal": 0.008056640625, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e5m4fnu": {"min": 0, "max": 126976.0, "num_bits": 9, "sign": 0, "exponent": 5, "mantissa": 4, "min_normal": 3.24249267578125e-05, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e1m8fnu": {"min": 0, "max": 3.9921875, "num_bits": 9, "sign": 0, "exponent": 1, "mantissa": 8, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e2m7fnu": {"min": 0, "max": 7.96875, "num_bits": 9, "sign": 0, "exponent": 2, "mantissa": 7, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e3m6fnu": {"min": 0, "max": 31.75, "num_bits": 9, "sign": 0, "exponent": 3, "mantissa": 6, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e4m5fnu": {"min": 0, "max": 504.0, "num_bits": 9, "sign": 0, "exponent": 4, "mantissa": 5, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e5m4fnu": {"min": 0, "max": 126976.0, "num_bits": 9, "sign": 0, "exponent": 5, "mantissa": 4, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float8_e1m7fnu": {"min": 0, "max": 3.984375, "num_bits": 8, "sign": 0, "exponent": 1, "mantissa": 7, "min_normal": 1.0078125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e2m6fnu": {"min": 0, "max": 7.9375, "num_bits": 8, "sign": 0, "exponent": 2, "mantissa": 6, "min_normal": 0.5078125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e3m5fnu": {"min": 0, "max": 31.5, "num_bits": 8, "sign": 0, "exponent": 3, "mantissa": 5, "min_normal": 0.12890625, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e4m4fnu": {"min": 0, "max": 496.0, "num_bits": 8, "sign": 0, "exponent": 4, "mantissa": 4, "min_normal": 0.00830078125, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e5m3fnu": {"min": 0, "max": 122880.0, "num_bits": 8, "sign": 0, "exponent": 5, "mantissa": 3, "min_normal": 3.4332275390625e-05, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e1m7fnu": {"min": 0, "max": 3.984375, "num_bits": 8, "sign": 0, "exponent": 1, "mantissa": 7, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e2m6fnu": {"min": 0, "max": 7.9375, "num_bits": 8, "sign": 0, "exponent": 2, "mantissa": 6, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e3m5fnu": {"min": 0, "max": 31.5, "num_bits": 8, "sign": 0, "exponent": 3, "mantissa": 5, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e4m4fnu": {"min": 0, "max": 496.0, "num_bits": 8, "sign": 0, "exponent": 4, "mantissa": 4, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e5m3fnu": {"min": 0, "max": 122880.0, "num_bits": 8, "sign": 0, "exponent": 5, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float7_e1m6fnu": {"min": 0, "max": 3.96875, "num_bits": 7, "sign": 0, "exponent": 1, "mantissa": 6, "min_normal": 1.015625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e2m5fnu": {"min": 0, "max": 7.875, "num_bits": 7, "sign": 0, "exponent": 2, "mantissa": 5, "min_normal": 0.515625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e3m4fnu": {"min": 0, "max": 31.0, "num_bits": 7, "sign": 0, "exponent": 3, "mantissa": 4, "min_normal": 0.1328125, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e4m3fnu": {"min": 0, "max": 480.0, "num_bits": 7, "sign": 0, "exponent": 4, "mantissa": 3, "min_normal": 0.0087890625, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e5m2fnu": {"min": 0, "max": 114688.0, "num_bits": 7, "sign": 0, "exponent": 5, "mantissa": 2, "min_normal": 3.814697265625e-05, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e1m6fnu": {"min": 0, "max": 3.96875, "num_bits": 7, "sign": 0, "exponent": 1, "mantissa": 6, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e2m5fnu": {"min": 0, "max": 7.875, "num_bits": 7, "sign": 0, "exponent": 2, "mantissa": 5, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e3m4fnu": {"min": 0, "max": 31.0, "num_bits": 7, "sign": 0, "exponent": 3, "mantissa": 4, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e4m3fnu": {"min": 0, "max": 480.0, "num_bits": 7, "sign": 0, "exponent": 4, "mantissa": 3, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e5m2fnu": {"min": 0, "max": 114688.0, "num_bits": 7, "sign": 0, "exponent": 5, "mantissa": 2, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float6_e1m5fnu": {"min": 0, "max": 3.9375, "num_bits": 6, "sign": 0, "exponent": 1, "mantissa": 5, "min_normal": 1.03125, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e2m4fnu": {"min": 0, "max": 7.75, "num_bits": 6, "sign": 0, "exponent": 2, "mantissa": 4, "min_normal": 0.53125, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e3m3fnu": {"min": 0, "max": 30.0, "num_bits": 6, "sign": 0, "exponent": 3, "mantissa": 3, "min_normal": 0.140625, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e4m2fnu": {"min": 0, "max": 448.0, "num_bits": 6, "sign": 0, "exponent": 4, "mantissa": 2, "min_normal": 0.009765625, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e5m1fnu": {"min": 0, "max": 98304.0, "num_bits": 6, "sign": 0, "exponent": 5, "mantissa": 1, "min_normal": 4.57763671875e-05, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e1m5fnu": {"min": 0, "max": 3.9375, "num_bits": 6, "sign": 0, "exponent": 1, "mantissa": 5, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e2m4fnu": {"min": 0, "max": 7.75, "num_bits": 6, "sign": 0, "exponent": 2, "mantissa": 4, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e3m3fnu": {"min": 0, "max": 30.0, "num_bits": 6, "sign": 0, "exponent": 3, "mantissa": 3, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e4m2fnu": {"min": 0, "max": 448.0, "num_bits": 6, "sign": 0, "exponent": 4, "mantissa": 2, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e5m1fnu": {"min": 0, "max": 98304.0, "num_bits": 6, "sign": 0, "exponent": 5, "mantissa": 1, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float5_e1m4fnu": {"min": 0, "max": 3.875, "num_bits": 5, "sign": 0, "exponent": 1, "mantissa": 4, "min_normal": 1.0625, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e2m3fnu": {"min": 0, "max": 7.5, "num_bits": 5, "sign": 0, "exponent": 2, "mantissa": 3, "min_normal": 0.5625, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e3m2fnu": {"min": 0, "max": 28.0, "num_bits": 5, "sign": 0, "exponent": 3, "mantissa": 2, "min_normal": 0.15625, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e4m1fnu": {"min": 0, "max": 384.0, "num_bits": 5, "sign": 0, "exponent": 4, "mantissa": 1, "min_normal": 0.01171875, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e5m0fnu": {"min": 0, "max": 65536.0, "num_bits": 5, "sign": 0, "exponent": 5, "mantissa": 0, "min_normal": 6.103515625e-05, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e1m4fnu": {"min": 0, "max": 3.875, "num_bits": 5, "sign": 0, "exponent": 1, "mantissa": 4, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e2m3fnu": {"min": 0, "max": 7.5, "num_bits": 5, "sign": 0, "exponent": 2, "mantissa": 3, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e3m2fnu": {"min": 0, "max": 28.0, "num_bits": 5, "sign": 0, "exponent": 3, "mantissa": 2, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e4m1fnu": {"min": 0, "max": 384.0, "num_bits": 5, "sign": 0, "exponent": 4, "mantissa": 1, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e5m0fnu": {"min": 0, "max": 65536.0, "num_bits": 5, "sign": 0, "exponent": 5, "mantissa": 0, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float4_e1m3fnu": {"min": 0, "max": 3.75, "num_bits": 4, "sign": 0, "exponent": 1, "mantissa": 3, "min_normal": 1.125, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e2m2fnu": {"min": 0, "max": 7.0, "num_bits": 4, "sign": 0, "exponent": 2, "mantissa": 2, "min_normal": 0.625, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e3m1fnu": {"min": 0, "max": 24.0, "num_bits": 4, "sign": 0, "exponent": 3, "mantissa": 1, "min_normal": 0.1875, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e4m0fnu": {"min": 0, "max": 256.0, "num_bits": 4, "sign": 0, "exponent": 4, "mantissa": 0, "min_normal": 0.015625, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e1m3fnu": {"min": 0, "max": 3.75, "num_bits": 4, "sign": 0, "exponent": 1, "mantissa": 3, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e2m2fnu": {"min": 0, "max": 7.0, "num_bits": 4, "sign": 0, "exponent": 2, "mantissa": 2, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e3m1fnu": {"min": 0, "max": 24.0, "num_bits": 4, "sign": 0, "exponent": 3, "mantissa": 1, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e4m0fnu": {"min": 0, "max": 256.0, "num_bits": 4, "sign": 0, "exponent": 4, "mantissa": 0, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float3_e1m2fnu": {"min": 0, "max": 3.5, "num_bits": 3, "sign": 0, "exponent": 1, "mantissa": 2, "min_normal": 1.25, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e2m1fnu": {"min": 0, "max": 6.0, "num_bits": 3, "sign": 0, "exponent": 2, "mantissa": 1, "min_normal": 0.75, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e3m0fnu": {"min": 0, "max": 16.0, "num_bits": 3, "sign": 0, "exponent": 3, "mantissa": 0, "min_normal": 0.25, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e1m2fnu": {"min": 0, "max": 3.5, "num_bits": 3, "sign": 0, "exponent": 1, "mantissa": 2, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e2m1fnu": {"min": 0, "max": 6.0, "num_bits": 3, "sign": 0, "exponent": 2, "mantissa": 1, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e3m0fnu": {"min": 0, "max": 16.0, "num_bits": 3, "sign": 0, "exponent": 3, "mantissa": 0, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float2_e1m1fnu": {"min": 0, "max": 3.0, "num_bits": 2, "sign": 0, "exponent": 1, "mantissa": 1, "min_normal": 1.5, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float2_e2m0fnu": {"min": 0, "max": 4.0, "num_bits": 2, "sign": 0, "exponent": 2, "mantissa": 0, "min_normal": 1.0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float2_e1m1fnu": {"min": 0, "max": 3.0, "num_bits": 2, "sign": 0, "exponent": 1, "mantissa": 1, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float2_e2m0fnu": {"min": 0, "max": 4.0, "num_bits": 2, "sign": 0, "exponent": 2, "mantissa": 0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float1_e1m0fnu": {"min": 0, "max": 2.0, "num_bits": 1, "sign": 0, "exponent": 1, "mantissa": 0, "min_normal": 2.0, "target_dtype": "fp1", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float1_e1m0fnu": {"min": 0, "max": 2.0, "num_bits": 1, "sign": 0, "exponent": 1, "mantissa": 0, "target_dtype": "fp1", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
}
dtype_dict["fp32"] = dtype_dict["float32"]
@@ -335,98 +335,12 @@ weights_dtype_order = [
"uint16", "float16_e1m15fnu", "float16_e2m14fnu", "float16_e3m13fnu", "float16_e4m12fnu", "float16_e5m11fnu",
]
use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
def check_torch_compile() -> bool: # dynamo can be disabled after startup
return use_torch_compile and not torch._dynamo.config.disable # pylint: disable=protected-access
if devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_rdna2_and_older = bool(gfx_version < 0x940 or (gfx_version < 0x1100 and gfx_version >= 0x1000))
else:
is_rdna2_and_older = False
if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
if devices.backend == "cuda":
is_fp8_mm_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9))
elif devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_fp8_mm_supported = bool(gfx_version >= 0x1200 or (gfx_version >= 0x940 and gfx_version < 0x1000))
else:
is_fp8_mm_supported = False
else:
is_fp8_mm_supported = os.environ.get("SDNQ_ALLOW_FP8_MM", "0").lower() not in {"0", "false", "no"}
if os.environ.get("SDNQ_USE_OPENVINO_MM", None) is None:
use_openvino_mm = bool(devices.backend in {"cpu", "openvino"})
else:
use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
use_triton_mm = bool(is_rdna2_and_older or devices.backend == "zluda")
else:
use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None:
# row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
else:
use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
fp_mm_func = None
int_mm_func = None
if use_openvino_mm:
try:
from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
int_mm_func = openvino_int_mm
fp_mm_func = openvino_fp_mm
except Exception:
use_openvino_mm = False
elif use_triton_mm:
try:
from .kernels.triton_mm import triton_int_mm, triton_fp_mm
int_mm_func = triton_int_mm
fp_mm_func = triton_fp_mm
except Exception:
use_triton_mm = False
if fp_mm_func is None and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
try:
from .triton_mm import triton_fp_mm
fp_mm_func = triton_fp_mm
except Exception:
use_triton_mm = False
if int_mm_func is None:
int_mm_func = torch._int_mm
if fp_mm_func is None:
if devices.backend == "cuda":
def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
return torch.mm(x,y, out_dtype=torch.float32)
else:
def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
if y.dtype == torch.float8_e4m3fn:
fp16_scale = 4 * y.shape[-2]
else:
fp16_scale = 65536 * y.shape[-2]
in_scale = fp16_scale**0.5
x = x.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
y = y.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
return torch.mm(x,y).to(dtype=torch.float32).mul_(fp16_scale)
fp_mm_func = fp_mm_torch
if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
use_contiguous_int8_mm = bool(use_openvino_mm or is_rdna2_and_older or devices.backend in {"ipex", "mps", "openvino", "zluda"})
use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
else:
use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
use_contiguous_fp16_mm = use_contiguous_int8_mm
if use_torch_compile:
if hasattr(torch._dynamo.config, "recompile_limit"):
torch._dynamo.config.recompile_limit = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
modules/sdnq/dequantizer.py
+18 -7
View File
@@ -5,13 +5,18 @@ from dataclasses import dataclass
import torch
from modules import devices
from .common import dtype_dict, compile_func, use_contiguous_int8_mm, use_contiguous_fp16_mm, use_tensorwise_fp8_matmul
from .common import dtype_dict, compile_func
from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_tensorwise_fp8_matmul, is_fp8_compile_supported
from .quant_utils import quantize_int_mm, quantize_uint_mm, quantize_fp_mm, rotate_hadamard, get_hadamard
from .packed_int import unpack_int
from .packed_float import unpack_float
from .layers import SDNQLayer
def skip_fp8_compile(weights_dtype: str) -> bool: # triton has no e4m3 conversions before sm_89, compiled dequant would crash
return not is_fp8_compile_supported and dtype_dict[weights_dtype]["storage_dtype"] == torch.float8_e4m3fn
@devices.inference_context()
def dequantize_asymmetric(
weight: torch.Tensor,
@@ -118,7 +123,7 @@ def re_quantize_int_mm(weight: torch.FloatTensor, matmul_dtype: str = "int8") ->
weight, scale = quantize_int_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale = quantize_int_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale = weight.t_(), scale.t_()
weight, scale = weight.t_(), scale.t_().contiguous()
return weight, scale
@@ -130,7 +135,7 @@ def re_quantize_uint_mm(weight: torch.FloatTensor, matmul_dtype: str = "uint8")
weight, scale, zero_point = quantize_uint_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale, zero_point = quantize_uint_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale, zero_point = weight.t_(), scale.t_(), zero_point.t_()
weight, scale, zero_point = weight.t_(), scale.t_().contiguous(), zero_point.t_().contiguous()
return weight, scale, zero_point
@@ -142,7 +147,7 @@ def re_quantize_fp_mm(weight: torch.FloatTensor, matmul_dtype: str = "float8_e4m
weight, scale = quantize_fp_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale = quantize_fp_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale = weight.t_(), scale.t_()
weight, scale = weight.t_(), scale.t_().contiguous()
if not use_tensorwise_fp8_matmul and dtype_dict[matmul_dtype]["num_bits"] == 8:
scale = scale.to(dtype=torch.float32)
return weight, scale
@@ -272,10 +277,14 @@ class SDNQDequantizer:
self.use_stochastic_rounding = use_stochastic_rounding
self.use_hadamard = use_hadamard
self.layer_class_name = layer_class_name
self.num_bits = dtype_dict[weights_dtype]["num_bits"]
self.is_packed = dtype_dict[weights_dtype]["is_packed"]
self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"]
self.is_integer = dtype_dict[weights_dtype]["is_integer"]
self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"]
self.num_bits_matmul = dtype_dict[quantized_matmul_dtype]["num_bits"]
self.is_packed_matmul = dtype_dict[quantized_matmul_dtype]["is_packed"]
self.is_integer_matmul = dtype_dict[quantized_matmul_dtype]["is_integer"]
self.is_unsigned_matmul = dtype_dict[quantized_matmul_dtype]["is_unsigned"]
@devices.inference_context()
def re_quantize_matmul(
@@ -287,10 +296,12 @@ class SDNQDequantizer:
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
non_hadamard: bool = True,
skip_compile: bool = False,
) -> tuple[torch.Tensor, torch.FloatTensor]: # pylint: disable=unused-argument
if hadamard is None and self.use_hadamard and not non_hadamard:
hadamard = get_hadamard(self.hadamard_group_size, dtype=self.result_dtype, device=weight.device)
return re_quantize_matmul_compiled(
re_quantize_matmul_func = re_quantize_matmul if skip_compile or skip_fp8_compile(self.weights_dtype) else re_quantize_matmul_compiled
return re_quantize_matmul_func(
self.weights_dtype,
weight,
scale,
@@ -322,7 +333,7 @@ class SDNQDequantizer:
if hadamard is None and self.use_hadamard and not non_hadamard:
hadamard = get_hadamard(self.hadamard_group_size, dtype=dtype, device=weight.device)
re_quantize_for_matmul = self.re_quantize_for_matmul or self.is_packed
dequantize_weight_func = dequantize_weight if skip_compile else dequantize_weight_compiled
dequantize_weight_func = dequantize_weight if skip_compile or skip_fp8_compile(self.weights_dtype) else dequantize_weight_compiled
return dequantize_weight_func(
self.weights_dtype,
weight,
modules/sdnq/forward.py
+5 -13
View File
@@ -2,7 +2,7 @@
from collections.abc import Callable
from .common import dtype_dict, embedding_types, conv_types, conv_transpose_types, use_tensorwise_fp8_matmul
from .common import dtype_dict, embedding_types, conv_types, conv_transpose_types
def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_quantized_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
@@ -20,12 +20,8 @@ def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_qua
return quantized_conv_forward_int8_matmul
else:
if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
if use_tensorwise_fp8_matmul:
from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
return quantized_conv_forward_fp8_matmul
else:
from .layers.conv.conv_fp8_scaled import quantized_conv_forward_fp8_scaled_matmul
return quantized_conv_forward_fp8_scaled_matmul
from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
return quantized_conv_forward_fp8_matmul
else:
from .layers.conv.conv_fp16 import quantized_conv_forward_fp16_matmul
return quantized_conv_forward_fp16_matmul
@@ -53,12 +49,8 @@ def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_qua
return quantized_linear_forward_int8_matmul
else:
if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
if use_tensorwise_fp8_matmul:
from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
return quantized_linear_forward_fp8_matmul
else:
from .layers.linear.linear_fp8_scaled import quantized_linear_forward_fp8_scaled_matmul
return quantized_linear_forward_fp8_scaled_matmul
from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
return quantized_linear_forward_fp8_matmul
else:
from .layers.linear.linear_fp16 import quantized_linear_forward_fp16_matmul
return quantized_linear_forward_fp16_matmul
modules/sdnq/kernel_wrappers.py
+155
View File
@@ -0,0 +1,155 @@
# pylint: disable=protected-access
import os
import sys
import torch
from modules import devices
if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
if devices.backend == "cuda":
is_fp8_mm_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9))
elif devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_fp8_mm_supported = bool(gfx_version >= 0x1200 or (gfx_version >= 0x940 and gfx_version < 0x1000))
else:
is_fp8_mm_supported = False
else:
is_fp8_mm_supported = os.environ.get("SDNQ_ALLOW_FP8_MM", "0").lower() not in {"0", "false", "no"}
if os.environ.get("SDNQ_ALLOW_FP8_COMPILE", None) is None:
if devices.backend == "cuda" and "linux" in sys.platform:
is_fp8_compile_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9)) # triton has no e4m3 conversions before sm_89
else:
is_fp8_compile_supported = True
else:
is_fp8_compile_supported = os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"}
if devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_rdna2_and_older = bool(gfx_version < 0x940 or (gfx_version < 0x1100 and gfx_version >= 0x1000))
else:
is_rdna2_and_older = False
if os.environ.get("SDNQ_USE_OPENVINO_MM", None) is None:
use_openvino_mm = bool(devices.backend in {"cpu", "openvino"})
else:
use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
use_triton_mm = bool(is_rdna2_and_older or devices.backend in {"zluda", "ipex", "xpu"})
else:
use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None:
# row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
else:
use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
use_contiguous_int8_mm = bool(use_openvino_mm or is_rdna2_and_older or devices.backend in {"ipex", "xpu", "mps", "openvino", "zluda"})
use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
else:
use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
use_contiguous_fp16_mm = use_contiguous_int8_mm
int_mm_func = None
fp_mm_func = None
fp8_mm_func = None
int_scaled_mm_func = None
fp_scaled_mm_func = None
fp8_scaled_mm_func = None
if use_openvino_mm:
try:
from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
int_mm_func = openvino_int_mm
fp_mm_func = openvino_fp_mm
except Exception:
use_openvino_mm = False
elif use_triton_mm:
try:
from .kernels.triton_mm import sdnq_triton_mm
from .kernels.triton_scaled_mm import sdnq_scaled_mm
int_mm_func = sdnq_triton_mm
fp_mm_func = sdnq_triton_mm
int_scaled_mm_func = sdnq_scaled_mm
fp_scaled_mm_func = sdnq_scaled_mm
if is_fp8_mm_supported:
fp8_mm_func = sdnq_triton_mm
fp8_scaled_mm_func = sdnq_scaled_mm
use_tensorwise_fp8_matmul = False
except Exception:
use_triton_mm = False
if fp_mm_func is None and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
try:
from .kernels.triton_mm import sdnq_triton_mm
from .kernels.triton_scaled_mm import sdnq_scaled_mm
fp_mm_func = sdnq_triton_mm
fp_scaled_mm_func = sdnq_scaled_mm
except Exception:
use_triton_mm = False
if int_mm_func is None:
def int_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
return torch._int_mm(a,b).to(dtype=out_dtype)
int_mm_func = int_mm_torch
if fp_mm_func is None:
if devices.backend == "cuda":
def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
return torch.mm(a,b, out_dtype=out_dtype)
else:
def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if b.dtype == torch.float8_e4m3fn:
fp16_scale = 4 * b.shape[-2]
else:
fp16_scale = 65536 * b.shape[-2]
in_scale = fp16_scale**0.5
a = a.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
b = b.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
return torch.mm(a,b).to(dtype=torch.float32).mul_(fp16_scale).to(dtype=out_dtype)
fp_mm_func = fp_mm_torch
if fp8_mm_func is None:
def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
dummy_input_scale = torch.ones(1, device=a.device, dtype=torch.float32)
return torch._scaled_mm(a, b, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=out_dtype)
fp8_mm_func = fp8_mm_torch
if int_scaled_mm_func is None:
def int_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return int_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, int_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
int_scaled_mm_func = int_scaled_mm_torch
if fp_scaled_mm_func is None:
def fp_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return fp_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, fp_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
fp_scaled_mm_func = fp_scaled_mm_torch
if fp8_scaled_mm_func is None:
if use_tensorwise_fp8_matmul:
def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
else:
def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is not None and bias.ndim != 1:
return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=None, out_dtype=out_dtype).add_(bias)
else:
return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=bias.to(dtype=out_dtype) if bias is not None else None, out_dtype=out_dtype)
fp8_scaled_mm_func = fp8_scaled_mm_torch
modules/sdnq/kernels/openvino_mm.py
+12 -12
View File
@@ -15,25 +15,25 @@ for ov_device in core.get_available_devices():
core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor:
def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True))
infer_request.infer()
C = C.to(A.device)
C = C.to(A.device, dtype=out_dtype)
return C
@torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
if "GPU" not in OV_DEVICE:
cache_key = (OV_DEVICE, "int8", Tensor_A.shape, Tensor_B.shape)
else:
cache_key = (OV_DEVICE, "int8", None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
if "GPU" not in OV_DEVICE:
shape_a = ov.Shape(Tensor_A.shape)
@@ -75,15 +75,15 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Ten
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
@openvino_int_mm.register_fake
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)
@torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=())
def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8"
if mm_dtype == "fp8":
Tensor_A = Tensor_A.to(dtype=torch.float16)
@@ -94,7 +94,7 @@ def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tens
cache_key = (OV_DEVICE, mm_dtype, None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
if "GPU" not in OV_DEVICE:
shape_a = ov.Shape(Tensor_A.shape)
@@ -135,8 +135,8 @@ def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tens
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
@openvino_fp_mm.register_fake
def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)
modules/sdnq/kernels/triton_atten.py
+23 -17
View File
@@ -5,7 +5,8 @@ import triton
import triton.language as tl
from ..common import compile_func # pylint: disable=relative-beyond-top-level
from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, apply_hadamard # pylint: disable=relative-beyond-top-level
from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, get_hadamard_group_size, apply_hadamard # pylint: disable=relative-beyond-top-level
from ..utils import is_pow2, next_power_of_2 # pylint: disable=relative-beyond-top-level
min_block_size = int(os.environ.get("SDNQ_TRITON_ATTEN_MIN_BLOCK_SIZE", "256"))
@@ -173,15 +174,15 @@ def sdnq_attn_kernel(
p_scale = tl.max(p, 1)[:, None] * (1 / 127.0)
p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale)
p = tl.floor(p * (1 / p_scale) + 0.5).to(tl.int8)
acc += tl.dot(p, v, out_dtype=tl.int32).to(tl.float32) * p_scale
acc = tl.fma(tl.dot(p, v, out_dtype=tl.int32).to(tl.float32), p_scale, acc)
else:
p_scale = tl.max(p, 1)[:, None] * (1 / (65504.0 if v.dtype == tl.float16 else 448.0))
p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale)
p = (p * (1 / p_scale)).to(v.dtype)
acc += tl.dot(p, v, out_dtype=tl.float32) * p_scale
acc = tl.fma(tl.dot(p, v, out_dtype=tl.float32), p_scale, acc)
else:
p = p.to(v.dtype)
acc += tl.dot(p, v, out_dtype=tl.float32)
acc = tl.dot(p, v, acc, out_dtype=tl.float32)
m_i = m_ij
l_i = 1 / l_i[:, None]
@@ -261,15 +262,15 @@ def get_attn_inputs(
out_dtype = query.dtype
if scale is None:
scale = QHD ** -0.5
if not math.log(QHD, 2).is_integer():
query = torch.nn.functional.pad(query, (0, triton.next_power_of_2(QHD) - QHD))
key = torch.nn.functional.pad(key, (0, triton.next_power_of_2(KHD) - KHD))
value = torch.nn.functional.pad(value, (0, triton.next_power_of_2(VHD) - VHD))
if not is_pow2(QHD):
query = torch.nn.functional.pad(query, (0, next_power_of_2(QHD) - QHD))
key = torch.nn.functional.pad(key, (0, next_power_of_2(KHD) - KHD))
value = torch.nn.functional.pad(value, (0, next_power_of_2(VHD) - VHD))
if attn_mask is not None:
attn_mask = attn_mask.expand((QZ, QH, QN, KN))
if not math.log(KN, 2).is_integer():
if not is_pow2(KN):
pad_value = float("-inf") if torch.is_floating_point(attn_mask) else 0
attn_mask = torch.nn.functional.pad(attn_mask, (0, triton.next_power_of_2(KN) - KN), value=pad_value)
attn_mask = torch.nn.functional.pad(attn_mask, (0, next_power_of_2(KN) - KN), value=pad_value)
if attn_mask.dtype == torch.bool:
attn_mask = attn_mask.to(dtype=torch.int8)
attn_mask = attn_mask.contiguous()
@@ -302,14 +303,18 @@ def sdnq_triton_atten(
do_quantize: bool = True,
out_dtype: torch.dtype | None = None,
) -> torch.FloatTensor:
QZ, QH, QN, _ = query.shape
_, _, KN, _ = key.shape
QZ, QH, QN, QHD = query.shape
_, _, KN, KHD = key.shape
_, _, VN, VHD = value.shape
if use_hadamard and matmul_dtype not in {None, "none", "no"}:
hadamard = get_hadamard(min(hadamard_group_size, query.shape[-1], key.shape[-1]), dtype=query.dtype, device=query.device)
else:
hadamard = None
hadamard = None
if use_hadamard and do_quantize and matmul_dtype not in {None, "none", "no"}:
hadamard_channel_size = next_power_of_2(min(QHD, KHD))
hadamard_group_size = min(hadamard_group_size, hadamard_channel_size)
use_hadamard, hadamard_group_size = get_hadamard_group_size(hadamard_channel_size, hadamard_group_size)
if use_hadamard:
hadamard = get_hadamard(hadamard_group_size, dtype=query.dtype, device=query.device)
(
query, query_scale,
key, key_scale,
@@ -324,6 +329,7 @@ def sdnq_triton_atten(
matmul_dtype=matmul_dtype, pv_matmul_dtype=pv_matmul_dtype,
do_quantize=do_quantize, out_dtype=out_dtype,
)
def grid(META):
return (triton.cdiv(QN, META["BLOCK_SIZE_M"]), QH, QZ)
out = torch.empty((QZ, QH, QN, value.shape[-1]), dtype=out_dtype, device=query.device)
@@ -346,4 +352,4 @@ def sdnq_triton_atten(
return out[..., :VHD]
get_attn_inputs = compile_func(get_attn_inputs, dynamic=True)
get_attn_inputs = compile_func(get_attn_inputs)
modules/sdnq/kernels/triton_mm.py
+61 -145
View File
@@ -1,12 +1,3 @@
"""
Modified from Triton MatMul example.
PyTorch torch._int_mm is broken on backward pass with Nvidia, so we use Triton on the backward pass with Nvidia.
AMD RDNA2 doesn't support torch._int_mm as it requires INT8 WMMA, so we use INT8 DP4A via Triton.
PyTorch doesn't support FP32 output type with FP16 MM, so we use Triton for FP16 MM too.
matmul_configs we use takes AMD and Intel into consideration too.
SDNQ Triton configs can outperform RocBLAS and OneDNN.
"""
import os
import math
import torch
@@ -18,92 +9,24 @@ import triton.language as tl
min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "256"))
matmul_configs = [
triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64").replace(" ","").split(",")]
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "2,4,8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "4,8").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "2").replace(" ","").split(",")]
]
@triton.autotune(configs=matmul_configs, key=["M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.jit
def triton_mm_kernel(
a_ptr, b_ptr, c_ptr,
M: tl.constexpr,
N: tl.constexpr,
K: tl.constexpr,
M_AT: tl.constexpr, # pylint: disable=unused-argument
N_AT: tl.constexpr, # pylint: disable=unused-argument
K_AT: tl.constexpr, # pylint: disable=unused-argument
stride_am: tl.constexpr, stride_ak: tl.constexpr,
stride_bk: tl.constexpr, stride_bn: tl.constexpr,
stride_cm: tl.constexpr, stride_cn: tl.constexpr,
a_dtype: tl.constexpr, # pylint: disable=unused-argument
out_dtype: tl.constexpr, # pylint: disable=unused-argument
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
) -> None:
pid = tl.program_id(axis=0)
num_pid_m: tl.constexpr = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group: tl.constexpr = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
tl.assume(M > 0)
tl.assume(N > 0)
tl.assume(K > 0)
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(stride_am > 0)
tl.assume(stride_ak > 0)
tl.assume(stride_bn > 0)
tl.assume(stride_bk > 0)
tl.assume(stride_cm > 0)
tl.assume(stride_cn > 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(BLOCK_SIZE_K > 0)
tl.assume(GROUP_SIZE_M > 0)
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
accumulator_dtype = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
accumulator = accumulator.to(c_ptr.type.element_ty)
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, accumulator, mask=c_mask)
# Intel requires tensor descriptors to perform good
@triton.autotune(configs=matmul_configs, key=["M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.jit
def triton_mm_td_kernel(
a_ptr, b_ptr, c_ptr,
def sdnq_triton_mm_kernel(
a_ptr, b_ptr, c_ptr, bias_ptr,
M: tl.constexpr,
N: tl.constexpr,
K: tl.constexpr,
b_is_contiguous: tl.constexpr,
bias_ndim: tl.constexpr,
M_AT: tl.constexpr, # pylint: disable=unused-argument
N_AT: tl.constexpr, # pylint: disable=unused-argument
K_AT: tl.constexpr, # pylint: disable=unused-argument
@@ -123,93 +46,86 @@ def triton_mm_td_kernel(
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
off_m = pid_m * BLOCK_SIZE_M
off_n = pid_n * BLOCK_SIZE_N
tl.assume(M > 0)
tl.assume(N > 0)
tl.assume(K > 0)
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(off_m >= 0)
tl.assume(off_n >= 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(BLOCK_SIZE_K > 0)
tl.assume(GROUP_SIZE_M > 0)
tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
if b_is_contiguous:
b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
else:
offs_k = tl.arange(0, BLOCK_SIZE_K)
offs_bn = (off_n + tl.arange(0, BLOCK_SIZE_N)) % N
b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
off_k = 0
accumulator_dtype = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator_dtype: tl.constexpr = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = a_desc.load([pid_m * BLOCK_SIZE_M, off_k])
b = b_desc.load([off_k, pid_n * BLOCK_SIZE_N])
a = a_desc.load([off_m, off_k])
if b_is_contiguous:
b = b_desc.load([off_k, off_n])
else:
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - off_k, other=0.0)
b_ptrs += BLOCK_SIZE_K
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
off_k += BLOCK_SIZE_K
if bias_ndim == 1:
accumulator = accumulator.to(tl.float32)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
bias = bias_desc.load([off_n])[None, :].to(tl.float32)
accumulator += bias
elif bias_ndim == 2:
accumulator = accumulator.to(tl.float32)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
bias = bias_desc.load([off_m, off_n]).to(tl.float32)
accumulator += bias
accumulator = accumulator.to(c_ptr.type.element_ty)
c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], accumulator)
c_desc.store([off_m, off_n], accumulator)
def triton_int_mm(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.Tensor:
def sdnq_triton_mm(
a: torch.Tensor,
b: torch.Tensor,
bias: torch.FloatTensor | None = None,
out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
if bias is not None:
assert bias.is_contiguous(), "Bias must be contiguous"
assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
M, K = a.shape
K, N = b.shape
if out_dtype is None:
out_dtype = torch.int32 if a.dtype == torch.int8 else torch.float32
c = torch.empty((M, N), device=a.device, dtype=out_dtype)
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
if b.is_contiguous():
triton_mm_td_kernel[grid](
a, b, c,
M, N, K,
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
else:
triton_mm_kernel[grid](
a, b, c,
M, N, K,
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
a.stride(0), a.stride(1),
b.stride(0), b.stride(1),
c.stride(0), c.stride(1),
str(a.dtype), str(c.dtype),
)
return c
def triton_fp_mm(a: torch.FloatTensor, b: torch.FloatTensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
M, K = a.shape
K, N = b.shape
c = torch.empty((M, N), device=a.device, dtype=out_dtype)
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
if b.is_contiguous():
triton_mm_td_kernel[grid](
a, b, c,
M, N, K,
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
else:
triton_mm_kernel[grid](
a, b, c,
M, N, K,
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
a.stride(0), a.stride(1),
b.stride(0), b.stride(1),
c.stride(0), c.stride(1),
str(a.dtype), str(c.dtype),
)
sdnq_triton_mm_kernel[grid](
a, b, c, bias,
M, N, K,
(1 if b.is_contiguous() else 0),
(0 if bias is None else bias.ndim),
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
return c
modules/sdnq/kernels/triton_scaled_mm.py
+147
View File
@@ -0,0 +1,147 @@
"""
W4A8 fallback with Triton.
This is intended as a template for future INT4 MM kernels as Triton has no support for INT4 hardware yet.
"""
import os
import math
import torch
import triton
import triton.language as tl
min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "256"))
matmul_configs = [
triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "2").replace(" ","").split(",")]
]
@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.jit
def sdnq_scaled_mm_kernel(
a_ptr, b_ptr, c_ptr, bias_ptr,
scale_a_ptr, scale_b_ptr,
M: tl.constexpr,
N: tl.constexpr,
K: tl.constexpr,
b_is_contiguous: tl.constexpr,
bias_ndim: tl.constexpr,
M_AT: tl.constexpr, # pylint: disable=unused-argument
N_AT: tl.constexpr, # pylint: disable=unused-argument
K_AT: tl.constexpr, # pylint: disable=unused-argument
a_dtype: tl.constexpr, # pylint: disable=unused-argument
out_dtype: tl.constexpr, # pylint: disable=unused-argument
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
) -> None:
pid = tl.program_id(axis=0)
num_pid_m: tl.constexpr = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group: tl.constexpr = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
off_m = pid_m * BLOCK_SIZE_M
off_n = pid_n * BLOCK_SIZE_N
tl.assume(M > 0)
tl.assume(N > 0)
tl.assume(K > 0)
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(off_m >= 0)
tl.assume(off_n >= 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(BLOCK_SIZE_K > 0)
tl.assume(GROUP_SIZE_M > 0)
tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
if b_is_contiguous:
b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
else:
offs_k = tl.arange(0, BLOCK_SIZE_K)
offs_bn = (off_n + tl.arange(0, BLOCK_SIZE_N)) % N
b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
off_k = 0
accumulator_dtype: tl.constexpr = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = a_desc.load([off_m, off_k])
if b_is_contiguous:
b = b_desc.load([off_k, off_n])
else:
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - off_k, other=0.0)
b_ptrs += BLOCK_SIZE_K
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
off_k += BLOCK_SIZE_K
scale_a_desc = tl.make_tensor_descriptor(base=scale_a_ptr, shape=(M,), strides=(1,), block_shape=(BLOCK_SIZE_M,))
scale_b_desc = tl.make_tensor_descriptor(base=scale_b_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
scale_a = scale_a_desc.load([off_m])[:, None].to(tl.float32)
scale_b = scale_b_desc.load([off_n])[None, :].to(tl.float32)
if bias_ndim == 1:
accumulator = accumulator.to(tl.float32) * scale_a
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
bias = bias_desc.load([off_n])[None, :].to(tl.float32)
accumulator = tl.fma(accumulator, scale_b, bias)
elif bias_ndim == 2:
accumulator = accumulator.to(tl.float32) * scale_a
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
bias = bias_desc.load([off_m, off_n]).to(tl.float32)
accumulator = tl.fma(accumulator, scale_b, bias)
else:
accumulator = accumulator.to(tl.float32) * scale_a * scale_b
accumulator = accumulator.to(c_ptr.type.element_ty)
c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
c_desc.store([off_m, off_n], accumulator)
def sdnq_scaled_mm(
a: torch.Tensor,
b: torch.Tensor,
scale_a: torch.Tensor,
scale_b: torch.Tensor,
bias: torch.FloatTensor | None = None,
out_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
assert scale_a.is_contiguous(), "Matrix A scale must be contiguous"
assert scale_b.is_contiguous(), "Matrix B scale must be contiguous"
if bias is not None:
assert bias.is_contiguous(), "Bias must be contiguous"
assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
M, K = a.shape
K, N = b.shape
c = torch.empty((M, N), device=a.device, dtype=out_dtype)
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
sdnq_scaled_mm_kernel[grid](
a, b, c, bias,
scale_a, scale_b,
M, N, K,
(1 if b.is_contiguous() else 0),
(0 if bias is None else bias.ndim),
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
return c
modules/sdnq/layers/conv/conv_fp16.py
+7 -6
View File
@@ -2,7 +2,8 @@
import torch
from ...common import compile_func, fp_mm_func
from ...common import compile_func
from ...kernel_wrappers import fp_mm_func, fp_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
@@ -50,7 +51,7 @@ def conv_fp16_matmul(
input, weight = check_mats(input, weight)
if groups == 1:
result = fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
result = fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -58,10 +59,10 @@ def conv_fp16_matmul(
for i in range(groups):
result.append(fp_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
modules/sdnq/layers/conv/conv_fp8.py
+8 -8
View File
@@ -3,6 +3,7 @@
import torch
from ...common import compile_func
from ...kernel_wrappers import fp8_mm_func, fp8_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
@@ -46,21 +47,20 @@ def conv_fp8_matmul(
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
input, weight = check_mats(input, weight)
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
if groups == 1:
result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale)
result = fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
for i in range(groups):
result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype))
result = torch.cat(result, dim=-1).mul_(input_scale)
if bias is not None:
dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
result.append(fp8_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
modules/sdnq/layers/conv/conv_fp8_scaled.py
-108
View File
@@ -1,108 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import get_conv_args, process_conv_input
from ..linear.linear_fp8_scaled import quantize_fp_scaled_mm_input
from ..linear.forward import check_mats
def conv_fp8_scaled_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
result_shape: torch.Size,
reversed_padding_repeated_twice: list[int],
padding_mode: str, conv_type: int,
groups: int, stride: list[int],
padding: list[int], dilation: list[int],
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
scale = scale.t()
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_scaled_mm_input(input)
input, weight = check_mats(input, weight)
if groups == 1:
if bias is not None and bias.dtype != torch.bfloat16:
bias = bias.to(dtype=torch.bfloat16)
result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
else:
scale = scale.view(groups, 1, scale.shape[1] // groups)
input_scale = input_scale.view(groups, input_scale.shape[0] // groups, 1)
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
if bias is not None:
bias = bias.view(groups, bias.shape[0] // groups)
if bias.dtype != torch.bfloat16:
bias = bias.to(dtype=torch.bfloat16)
for i in range(groups):
result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=torch.bfloat16))
else:
for i in range(groups):
result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=torch.bfloat16))
result = torch.cat(result, dim=-1)
if svd_up is not None:
result.add_(svd_bias)
result = result.view(mm_output_shape).to(return_dtype)
if conv_type == 1:
result = result.transpose_(1,2)
elif conv_type == 2:
result = result.permute(0,3,1,2)
elif conv_type == 3:
result = result.permute(0,4,1,2,3)
return result
def quantized_conv_forward_fp8_scaled_matmul(self, input) -> torch.FloatTensor:
if torch.numel(input) / input.shape[2] < 32:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp8_scaled_matmul(
input, weight, scale,
self.sdnq_dequantizer.result_shape,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
conv_fp8_scaled_matmul = compile_func(conv_fp8_scaled_matmul)
modules/sdnq/layers/conv/conv_int8.py
+7 -6
View File
@@ -2,7 +2,8 @@
import torch
from ...common import compile_func, int_mm_func
from ...common import compile_func
from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
@@ -64,7 +65,7 @@ def conv_int8_matmul(
input, weight = check_mats(input, weight)
if groups == 1:
result = int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
result = int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -72,10 +73,10 @@ def conv_int8_matmul(
for i in range(groups):
result.append(int_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
modules/sdnq/layers/conv/conv_uint8.py
+4 -3
View File
@@ -2,7 +2,8 @@
import torch
from ...common import compile_func, int_mm_func
from ...common import compile_func
from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
from ...dequantizer import dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
@@ -67,7 +68,7 @@ def conv_uint8_matmul(
input, weight = check_mats(input, weight)
if groups == 1:
result = int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
result = int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -75,7 +76,7 @@ def conv_uint8_matmul(
for i in range(groups):
result.append(int_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
result = dequantize_asymmetric(result, scale, zero_bias, dtype=return_dtype, result_shape=mm_output_shape)
result = dequantize_asymmetric(result, scale, zero_bias, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
modules/sdnq/layers/linear/forward.py
+1 -1
View File
@@ -2,7 +2,7 @@
import torch
from ...common import use_contiguous_int8_mm, use_contiguous_fp16_mm
from ...kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
def check_mats(input: torch.Tensor, weight: torch.Tensor, allow_contiguous_mm: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
modules/sdnq/layers/linear/linear_fp16.py
+5 -6
View File
@@ -2,8 +2,8 @@
import torch
from ...common import compile_func, fp_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...common import compile_func
from ...kernel_wrappers import fp_scaled_mm_func
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
@@ -29,6 +29,7 @@ def fp16_matmul(
weight = weight.to(dtype=torch.float16) # fp8 weights
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
@@ -37,12 +38,10 @@ def fp16_matmul(
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype, matmul_dtype="float16")
input, weight = check_mats(input, weight)
if bias is not None:
return dequantize_asymmetric(fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
else:
return dequantize_symmetric(fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
return fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
modules/sdnq/layers/linear/linear_fp8.py
+4 -6
View File
@@ -3,7 +3,7 @@
import torch
from ...common import compile_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...kernel_wrappers import fp8_scaled_mm_func
from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
@@ -36,6 +36,7 @@ def fp8_matmul(
scale = scale.t()
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
@@ -44,13 +45,10 @@ def fp8_matmul(
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
input, weight = check_mats(input, weight, allow_contiguous_mm=False)
if bias is not None:
return dequantize_asymmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
else:
return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
return fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
modules/sdnq/layers/linear/linear_fp8_scaled.py
-75
View File
@@ -1,75 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import check_mats
def quantize_fp_scaled_mm_input(input: torch.FloatTensor, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]:
input = input.flatten(0,-2).to(dtype=torch.float32)
input, input_scale = quantize_fp_mm(input, dim=-1, matmul_dtype=matmul_dtype)
return input, input_scale
def fp8_scaled_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
scale = scale.t()
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_scaled_mm_input(input)
input, weight = check_mats(input, weight, allow_contiguous_mm=False)
if bias is not None and bias.dtype != torch.bfloat16:
bias = bias.to(dtype=torch.bfloat16)
result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
if svd_up is not None:
result.add_(svd_bias)
result = result.view(output_shape).to(return_dtype)
return result
def quantized_linear_forward_fp8_scaled_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
return fp8_scaled_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
fp8_scaled_matmul = compile_func(fp8_scaled_matmul)
modules/sdnq/layers/linear/linear_int8.py
+6 -9
View File
@@ -2,19 +2,19 @@
import torch
from ...common import compile_func, int_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...common import compile_func
from ...kernel_wrappers import int_scaled_mm_func
from ...quant_utils import quantize_int_mm, rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import check_mats
def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.FloatTensor]:
def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.FloatTensor]:
input = input.flatten(0,-2)
if dtype is not None:
input = input.to(dtype=dtype)
input, input_scale = quantize_int_mm(input, dim=-1)
input, input_scale = quantize_int_mm(input, dim=-1, matmul_dtype=matmul_dtype)
if input_scale.dtype == torch.float16: # fp16 will overflow
input_scale = input_scale.to(dtype=torch.float32)
return input, input_scale
@@ -64,12 +64,9 @@ def int8_matmul(
if bias is not None:
zero_bias.add_(bias)
bias = zero_bias
input, weight = check_mats(input, weight)
if bias is not None:
return dequantize_asymmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
else:
return dequantize_symmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
input, weight = check_mats(input, weight)
return int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
modules/sdnq/layers/linear/linear_uint8.py
+8 -8
View File
@@ -2,19 +2,19 @@
import torch
from ...common import compile_func, int_mm_func
from ...dequantizer import dequantize_asymmetric
from ...common import compile_func
from ...kernel_wrappers import int_scaled_mm_func
from ...quant_utils import quantize_uint_mm, rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import check_mats
def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None, matmul_dtype: str = "uint8") -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
input = input.flatten(0,-2)
if dtype is not None:
input = input.to(dtype=dtype)
input, input_scale, input_zero_point = quantize_uint_mm(input, dim=-1)
input, input_scale, input_zero_point = quantize_uint_mm(input, dim=-1, matmul_dtype=matmul_dtype)
if input_scale.dtype == torch.float16: # fp16 will overflow
input_scale = input_scale.to(dtype=torch.float32)
input_zero_point = input_zero_point.to(dtype=torch.float32)
@@ -61,16 +61,16 @@ def uint8_matmul(
input, input_scale, input_zero_point = quantize_uint_mm_input(input, dtype=scale.dtype)
if zero_point is not None:
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(input_scale.dtype).mul_(input_scale).mul(zero_point)
zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point))
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point)
zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point))
zero_bias.add_(torch.mul(input_zero_point.mul_(input.shape[-1]), zero_point))
else:
zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point)
zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point)
if bias is not None:
zero_bias.add_(bias)
input, weight = check_mats(input, weight)
return dequantize_asymmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, zero_bias, dtype=return_dtype, result_shape=output_shape)
return int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
modules/sdnq/loader.py
+5 -4
View File
@@ -3,7 +3,8 @@ import json
import torch
from modules import shared
from .common import dtype_dict, is_fp8_mm_supported, use_tensorwise_fp8_matmul, check_torch_compile, linear_types
from .common import dtype_dict, check_torch_compile, linear_types
from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul
from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, sdnq_post_load_quant
from .quant_utils import prepare_weight_for_matmul, prepare_svd_for_matmul
from .utils import get_quant_args_from_config, check_param_name_in
@@ -286,10 +287,10 @@ def apply_sdnq_options_to_module(
if current_use_quantized_matmul is not None:
if current_use_quantized_matmul != module.sdnq_dequantizer.use_quantized_matmul:
if not module.sdnq_dequantizer.re_quantize_for_matmul and not dtype_dict[module.sdnq_dequantizer.weights_dtype]["is_packed"]:
module.scale.t_()
module.weight.t_()
module.scale.data = module.scale.t_().contiguous()
module.weight.data = module.weight.t_()
if module.zero_point is not None:
module.zero_point.t_()
module.zero_point.data = module.zero_point.t_().contiguous()
if current_use_quantized_matmul:
module.weight.data = prepare_weight_for_matmul(module.weight, matmul_dtype=module.sdnq_dequantizer.quantized_matmul_dtype)
else:
modules/sdnq/packed_float.py
+23 -3
View File
@@ -37,7 +37,6 @@ def pack_float(x: torch.FloatTensor, weights_dtype: str) -> torch.Tensor:
mantissa_mask = (1 << mantissa_difference) # pylint: disable=superfluous-parens
x = x.to(dtype=torch.float32).view(torch.int32)
x = torch.where(
torch.gt(
torch.bitwise_and(x, -(1 << (mantissa_difference-4)) & ~(-mantissa_mask)),
@@ -47,7 +46,22 @@ def pack_float(x: torch.FloatTensor, weights_dtype: str) -> torch.Tensor:
x,
)
x = torch.where(torch.lt(x.view(torch.float32).abs(), dtype_dict[weights_dtype]["min_normal"]), 0, x)
if exponent_bits < 8:
min_normal = 2.0 ** (2 - (1 << (exponent_bits - 1)))
x_f32_abs = x.view(dtype=torch.float32).abs()
is_subnormal = torch.lt(x_f32_abs, min_normal)
x = torch.where(
is_subnormal,
torch.bitwise_or(
torch.bitwise_and(x, -2147483648),
torch.bitwise_left_shift(
x_f32_abs.mul_((1 << mantissa_bits) / min_normal).round_().to(torch.int32),
mantissa_difference,
),
),
x,
)
del x_f32_abs, is_subnormal
x = torch.bitwise_right_shift(x, mantissa_difference)
x = torch.bitwise_and(
@@ -105,5 +119,11 @@ def unpack_float(x: torch.Tensor, weights_dtype: str, shape: torch.Size) -> torc
overflow_mask = (~(-(1 << (22 + exponent_bits))) | 1090519039)
x = torch.where(torch.bitwise_and(x, overflow_mask).to(dtype=torch.bool), x, 0)
x = x.view(torch.float32)
if exponent_bits < 8:
min_normal = 2.0 ** (2 - (1 << (exponent_bits - 1)))
x = torch.where(
torch.lt(x.abs(), min_normal),
torch.sign(x).mul_(-min_normal).add_(x, alpha=2.0),
x,
)
return x
modules/sdnq/quant_utils.py
+17 -15
View File
@@ -1,10 +1,11 @@
# pylint: disable=redefined-builtin
import math
import torch
from modules import devices
from .common import dtype_dict, use_contiguous_int8_mm, use_contiguous_fp16_mm, conv_types, conv_transpose_types
from .common import dtype_dict, conv_types, conv_transpose_types
from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
from .utils import is_pow2, is_pow4, next_power_of_2
@devices.inference_context()
@@ -100,12 +101,12 @@ def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.de
@devices.inference_context()
def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
if math.log(n, 4).is_integer():
if is_pow4(n):
return build_hadamard_n4(n, device=device, dtype=dtype)
elif math.log(n, 2).is_integer():
elif is_pow2(n):
return build_hadamard_n2(n, device=device, dtype=dtype)
else:
raise RuntimeError("Hadamard Group Size must be a power of 2.")
raise RuntimeError(f"Hadamard Group Size must be a power of 2 but got {n}.")
# 256x256 Hadamard matrix is just 256 KB at FP32
@@ -140,10 +141,18 @@ def rotate_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch
return result
def get_hadamard_group_size(channel_size: int, group_size: int) -> tuple[bool, int]:
group_size = next_power_of_2(min(channel_size, group_size))
if channel_size % group_size != 0:
while channel_size % group_size != 0:
group_size = group_size // 2
use_hadamard = group_size >= 4
return use_hadamard, group_size
@devices.inference_context()
def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.FloatTensor | None = None, layer_class_name: str | None = None) -> torch.Tensor:
def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.FloatTensor | None = None, layer_class_name: str | None = None) -> tuple[torch.Tensor, bool, int]:
is_conv = False
use_hadamard = True
if hadamard is not None:
group_size = hadamard.shape[-1]
if layer_class_name in conv_types or layer_class_name in conv_transpose_types:
@@ -151,14 +160,7 @@ def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.
channel_size = weight.shape[1]
else:
channel_size = weight.shape[-1]
group_size = min(group_size, channel_size)
if channel_size % group_size != 0:
hadamard_pow2 = int(math.log2(group_size))
while channel_size % group_size != 0:
hadamard_pow2 -= 1
group_size = 2 ** hadamard_pow2
if group_size < 4:
use_hadamard = False
use_hadamard, group_size = get_hadamard_group_size(channel_size, group_size)
if use_hadamard:
if hadamard is not None and group_size != hadamard.shape[-1]:
hadamard = None
modules/sdnq/quantizer.py
+59 -53
View File
@@ -14,30 +14,43 @@ from accelerate import init_empty_weights
from modules import devices, shared
from .common import (
sdnq_version,
sdnq_keys,
dtype_dict,
accepted_weight_dtypes,
accepted_matmul_dtypes,
allowed_types,
linear_types,
embedding_types,
conv_types,
conv_transpose_types,
weights_dtype_order,
is_fp8_mm_supported,
use_tensorwise_fp8_matmul,
check_torch_compile,
compile_func,
)
from .quant_utils import (
quantize_weight,
apply_svdquant,
apply_hadamard,
prepare_weight_for_matmul,
prepare_svd_for_matmul,
)
from .utils import (
check_param_name_in,
check_quant_is_allowed,
check_quantized_matmul_is_allowed,
get_quant_args_from_config,
get_quant_kwargs,
get_quantized_matmul_dtype,
add_module_skip_keys,
)
from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul
from .dequantizer import SDNQDequantizer, dequantize_sdnq_model
from .packed_int import pack_int
from .packed_float import pack_float
from .forward import get_forward_func
from .layers import get_sdnq_wrapper_class
from .quant_utils import quantize_weight, apply_svdquant, apply_hadamard, prepare_weight_for_matmul, prepare_svd_for_matmul
from .utils import check_param_name_in, get_quant_args_from_config, get_quant_kwargs, get_quantized_matmul_dtype, add_module_skip_keys
from .common import sdnq_version as current_sdnq_version
class QuantizationMethod(str, Enum):
@@ -101,7 +114,7 @@ def sdnq_quantize_layer_weight(
is_conv_type = True
reduction_axes = 1
output_channel_size, channel_size = weight.shape[:2]
use_quantized_matmul = use_quantized_matmul and channel_size >= 32 and output_channel_size >= 32 and output_channel_size % 16 == 0 and channel_size % 16 == 0
use_quantized_matmul = check_quantized_matmul_is_allowed(use_quantized_matmul, output_channel_size, channel_size)
if use_quantized_matmul and not re_quantize_for_matmul and not dtype_dict[weights_dtype]["is_packed"]:
result_shape = weight.shape
weight = weight.flatten(1,-1)
@@ -115,7 +128,7 @@ def sdnq_quantize_layer_weight(
is_linear_type = True
reduction_axes = -1
output_channel_size, channel_size = weight.shape
use_quantized_matmul = use_quantized_matmul and channel_size >= 32 and output_channel_size >= 32 and output_channel_size % 16 == 0 and channel_size % 16 == 0
use_quantized_matmul = check_quantized_matmul_is_allowed(use_quantized_matmul, output_channel_size, channel_size)
else:
if weight.ndim > 1:
output_channel_size, channel_size = weight.shape[-2:]
@@ -206,10 +219,10 @@ def sdnq_quantize_layer_weight(
weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype, dtype=(scale_dtype if cast_scale else None), use_stochastic_rounding=(use_stochastic_rounding and not skip_sr))
if transpose_weights:
scale.t_()
weight.t_()
scale = scale.t_().contiguous()
weight = weight.t_()
if zero_point is not None:
zero_point.t_()
zero_point = zero_point.t_().contiguous()
weight = prepare_weight_for_matmul(weight, matmul_dtype=quantized_matmul_dtype)
quantized_weight_shape = weight.shape
@@ -435,7 +448,7 @@ def sdnq_quantize_layer(layer: torch.nn.Module, quantization_config: "SDNQConfig
@devices.inference_context()
def apply_sdnq_to_module(model: torch.nn.Module, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, full_param_name: str = "") -> tuple[torch.nn.Module, "SDNQConfig"]: # pylint: disable=unused-argument
def apply_sdnq_to_module(model: torch.nn.Module, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, pre_quantized: bool = False, full_param_name: str = "") -> tuple[torch.nn.Module, "SDNQConfig"]: # pylint: disable=unused-argument
if not list(model.children()):
return model, quantization_config
for module_name, module in model.named_children():
@@ -445,21 +458,13 @@ def apply_sdnq_to_module(model: torch.nn.Module, quantization_config: "SDNQConfi
param_name = module_name
if hasattr(module, "weight") and module.weight is not None:
param_name = param_name + ".weight"
layer_class_name = module.__class__.__name__
param_in_modules_to_not_convert = check_param_name_in(param_name, quantization_config.modules_to_not_convert)
if (
layer_class_name in allowed_types
and module.weight.dtype in {torch.float64, torch.float32, torch.float16, torch.bfloat16}
and param_in_modules_to_not_convert is None
and not (layer_class_name in embedding_types and not quantization_config.quant_embedding)
and not ((layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quantization_config.quant_conv)
):
module, quantization_config = sdnq_quantize_layer(module, quantization_config, torch_dtype=torch_dtype, param_name=param_name)
setattr(model, module_name, module)
elif param_in_modules_to_not_convert is None:
quantization_config.modules_to_not_convert.append(param_name)
module, quantization_config = apply_sdnq_to_module(module, quantization_config, torch_dtype=torch_dtype, full_param_name=param_name)
if check_param_name_in(param_name, quantization_config.modules_to_not_convert) is None:
if check_quant_is_allowed(module.__class__.__name__, module.weight, quantization_config, pre_quantized=pre_quantized):
module, quantization_config = sdnq_quantize_layer(module, quantization_config, torch_dtype=torch_dtype, param_name=param_name)
setattr(model, module_name, module)
else:
quantization_config.modules_to_not_convert.append(param_name)
module, quantization_config = apply_sdnq_to_module(module, quantization_config, torch_dtype=torch_dtype, pre_quantized=pre_quantized, full_param_name=param_name)
setattr(model, module_name, module)
return model, quantization_config
@@ -485,6 +490,7 @@ def sdnq_post_load_quant(
dequantize_fp32: bool = True,
non_blocking: bool = False,
add_skip_keys:bool = True,
minimum_allowed_numel: int = 16384,
modules_to_not_convert: list[str] | None = None,
modules_to_not_use_matmul: list[str] | None = None,
modules_dtype_dict: dict[str, list[str]] | None = None,
@@ -523,6 +529,7 @@ def sdnq_post_load_quant(
dequantize_fp32=dequantize_fp32,
non_blocking=non_blocking,
add_skip_keys=add_skip_keys,
minimum_allowed_numel=minimum_allowed_numel,
modules_to_not_convert=modules_to_not_convert,
modules_to_not_use_matmul=modules_to_not_use_matmul,
modules_dtype_dict=modules_dtype_dict,
@@ -534,7 +541,7 @@ def sdnq_post_load_quant(
model, quantization_config = add_module_skip_keys(model, quantization_config)
model.eval()
model, quantization_config = apply_sdnq_to_module(model, quantization_config, torch_dtype=torch_dtype)
model, quantization_config = apply_sdnq_to_module(model, quantization_config, torch_dtype=torch_dtype, pre_quantized=pre_quantized)
model.quantization_config = quantization_config
if hasattr(model, "config"):
@@ -602,17 +609,11 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
return True
elif param_name.endswith(".weight"):
if check_param_name_in(param_name, self.quantization_config.modules_to_not_convert) is None:
layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__
if layer_class_name in allowed_types:
if layer_class_name in embedding_types:
if self.quantization_config.quant_embedding:
return True
elif layer_class_name in conv_types or layer_class_name in conv_transpose_types:
if self.quantization_config.quant_conv:
return True
else:
return True
self.quantization_config.modules_to_not_convert.append(param_name)
layer = get_module_from_name(model, param_name)[0]
if check_quant_is_allowed(layer.__class__.__name__, layer.weight, self.quantization_config, pre_quantized=self.pre_quantized):
return True
else:
self.quantization_config.modules_to_not_convert.append(param_name)
return False
@devices.inference_context()
@@ -865,10 +866,8 @@ class SDNQConfig(QuantizationConfigMixin):
Enabling this option will use non blocking ops when moving layers between the quantization device and the return device.
add_skip_keys (`bool`, *optional*, defaults to `True`):
Disabling this option won't add model specific keys to modules_to_not_convert, modules_to_not_use_matmul and modules_dtype_dict.
quantization_device (`torch.device`, *optional*, defaults to `None`):
Used to set which device will be used for the quantization calculation on model load.
return_device (`torch.device`, *optional*, defaults to `None`):
Used to set which device will the quantized weights be sent back to.
minimum_allowed_numel (`int`, *optional*, defaults to `16384`):
Layers that have less than `minimum_allowed_numel` elements in them will be skipped and added to `modules_to_not_convert`.
modules_to_not_convert (`list`, *optional*, default to `None`):
The list of modules to not quantize. Useful for quantizing models that explicitly require to have some
modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
@@ -880,6 +879,10 @@ class SDNQConfig(QuantizationConfigMixin):
modules_quant_config (`dict`, *optional*, default to `None`):
The dict of modules and a dict of quantization kwargs to use for that module.
Useful for quantizing some modules with a different quantization config.
quantization_device (`torch.device`, *optional*, defaults to `None`):
Used to set which device will be used for the quantization calculation on model load.
return_device (`torch.device`, *optional*, defaults to `None`):
Used to set which device will the quantized weights be sent back to.
is_training (`bool`, *optional*, defaults to `False`):
This option is auto enabled when using `sdnq.training.sdnq_training_post_load_quant`.
Enabling this option with SDNQQuantizer will convert the model to SDNQ Training format after quantization.
@@ -908,22 +911,20 @@ class SDNQConfig(QuantizationConfigMixin):
dequantize_fp32: bool = True,
non_blocking: bool = False,
add_skip_keys: bool = True,
quantization_device: torch.device | None = None,
return_device: torch.device | None = None,
minimum_allowed_numel: int = 16384,
modules_to_not_convert: list[str] | None = None,
modules_to_not_use_matmul: list[str] | None = None,
modules_dtype_dict: dict[str, list[str]] | None = None,
modules_quant_config: dict[str, dict] | None = None,
quantization_device: torch.device | None = None,
return_device: torch.device | None = None,
sdnq_version: str | None = None,
is_training: bool = False,
**kwargs,
):
self.weights_dtype = weights_dtype
self.quantized_matmul_dtype = quantized_matmul_dtype
self.is_training = is_training
if self.is_training:
self.quant_method = QuantizationMethod.SDNQ_TRAINING
else:
self.quant_method = QuantizationMethod.SDNQ
self.hadamard_group_size = hadamard_group_size
self.group_size = group_size
self.svd_rank = svd_rank
@@ -942,14 +943,19 @@ class SDNQConfig(QuantizationConfigMixin):
self.dequantize_fp32 = dequantize_fp32
self.non_blocking = non_blocking
self.add_skip_keys = add_skip_keys
self.quantization_device = quantization_device
self.return_device = return_device
self.minimum_allowed_numel = minimum_allowed_numel
self.modules_to_not_convert = modules_to_not_convert
self.modules_to_not_use_matmul = modules_to_not_use_matmul
self.modules_dtype_dict = modules_dtype_dict
self.modules_quant_config = modules_quant_config
self.quantization_device = quantization_device
self.return_device = return_device
self.sdnq_version = current_sdnq_version if sdnq_version is None else sdnq_version
self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]
self.sdnq_version = sdnq_version
if self.is_training:
self.quant_method = QuantizationMethod.SDNQ_TRAINING
else:
self.quant_method = QuantizationMethod.SDNQ
self.post_init()
def post_init(self) -> None:
modules/sdnq/utils.py
+41 -1
View File
@@ -1,7 +1,29 @@
import re
import torch
from .common import dtype_dict, common_skip_keys, module_skip_keys_dict, conv_types, conv_transpose_types
from .common import (
dtype_dict,
common_skip_keys,
module_skip_keys_dict,
allowed_types,
embedding_types,
conv_types,
conv_transpose_types,
)
def is_pow2(n: int) -> bool:
return (n & (n - 1)) == 0
def is_pow4(n: int) -> bool:
return is_pow2(n) and (n.bit_length() & 1 == 1)
def next_power_of_2(n: int) -> int:
if is_pow2(n):
return n
return 2 ** n.bit_length()
def check_param_name_in(param_name: str, param_list: list[str]) -> str:
@@ -21,6 +43,24 @@ def check_param_name_in(param_name: str, param_list: list[str]) -> str:
return None
def check_quant_is_allowed(layer_class_name: str, weight: torch.Tensor, quantization_config, pre_quantized: bool = False) -> bool:
return bool(
layer_class_name in allowed_types
and weight.dtype in {torch.float64, torch.float32, torch.float16, torch.bfloat16}
and not (layer_class_name in embedding_types and not quantization_config.quant_embedding)
and not ((layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quantization_config.quant_conv)
and (pre_quantized or weight.numel() >= quantization_config.minimum_allowed_numel)
)
def check_quantized_matmul_is_allowed(use_quantized_matmul: bool, output_channel_size: int, channel_size: int) -> bool:
return bool(
use_quantized_matmul
and output_channel_size >= 32 and channel_size >= 32
and output_channel_size % 16 == 0 and channel_size % 16 == 0
)
def get_quant_args_from_config(quantization_config: dict) -> dict:
from .quantizer import SDNQConfig
if isinstance(quantization_config, SDNQConfig):
modules/shared_items.py
+1 -1
View File
@@ -179,5 +179,5 @@ def get_repo(model):
return None
sdnq_quant_modes = ["int8", "uint8", "int6", "uint6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float8_e3m4fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn"]
sdnq_quant_modes = ["int8", "uint8", "int6", "uint6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float8_e3m4fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn", "int16", "uint16", "float16"]
sdnq_matmul_modes = ["auto", "int8", "uint8", "float8_e4m3fn", "float16"]
modules/ui_caption.py
+1 -1
View File
@@ -361,7 +361,7 @@ def create_ui():
with gr.Column(variant='compact', elem_id='caption_output'):
with gr.Row(elem_id='caption_output_prompt'):
prompt = gr.Textbox(label="Answer", lines=12, placeholder="ai generated image description")
prompt = gr.Textbox(label="Answer", lines=12, placeholder="VLM generated image description")
with gr.Row(elem_id='caption_output_image'):
output_image = gr.Image(type='pil', label="Annotated Image", interactive=False, visible=False, elem_id='caption_output_image_display')
with gr.Row(elem_id='caption_output_classes'):
modules/ui_common.py
+2
View File
@@ -180,6 +180,8 @@ def save_files(js_data, files, html_info, index):
if index < len(files):
files = [files[index]]
start_index = index
elif len(files) == 1:
start_index = 0
else:
log.error(f'Save: index={index} first={p.index_of_first_image} files={len(files)} out of range')
files = []
modules/ui_definitions.py
+1
View File
@@ -94,6 +94,7 @@ def create_settings(cmd_opts):
"sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_titles()}, refresh=refresh_checkpoints),
"sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles()}, refresh=refresh_checkpoints),
"sd_unet": OptionInfo("Default", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list),
"sd_unet_secondary": OptionInfo("Default", "UNET model secondary", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list),
"latent_history": OptionInfo(20, "Latent history size", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
"advanced_sep": OptionInfo("<h2>Advanced Options</h2>", "", gr.HTML),
modules/ui_extra_networks.py
+1 -1
View File
@@ -774,7 +774,7 @@ def create_ui(container, button_parent: gr.Button, tabname: str, skip_indexing =
def ui_tab_change(page):
scan_visible = page in ['Model', 'Lora', 'VAE', 'UNet/DiT', 'Hypernetwork', 'Embedding']
save_visible = page in ['Style']
model_visible = page in ['Model']
model_visible = page in ['Model', 'UNet/DiT']
return [gr.update(visible=scan_visible), gr.update(visible=save_visible), gr.update(visible=model_visible)]
ui.button_refresh = ui_components.ToolButton(ui_symbols.refresh, elem_id=f"{tabname}_extra_refresh")
modules/ui_settings.py
+9 -2
View File
@@ -202,7 +202,7 @@ def run_settings_single(value, key, progress=False, force=False):
shared.opts.save(silent=True)
if key == 'sd_text_encoder':
sd_models.reload_text_encoder() # apply the change now; reloads the model for encoders with no in-place swap
if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet'] or force:
if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet', 'sd_unet_secondary'] or force:
log.debug(f'Setting changed: {key}="{value}" progress={progress} force={force}')
return get_value_for_setting(key), shared.opts.dumpjson()
@@ -442,7 +442,7 @@ def create_quicksettings(interfaces):
)
button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False)
button_set_refiner.click(
fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'),
fn=lambda value, _: run_settings_single(value, key='sd_model_refiner'),
_js="consumeDesiredCheckpointName",
inputs=[shared.settings_components['sd_model_refiner'], dummy_component],
outputs=[shared.settings_components['sd_model_refiner'], text_settings],
@@ -461,6 +461,13 @@ def create_quicksettings(interfaces):
inputs=[shared.settings_components["sd_unet"], dummy_component],
outputs=[shared.settings_components["sd_unet"], text_settings],
)
button_set_unet_secondary = gr.Button("Change UNet secondary", elem_id="change_unet_secondary", visible=False)
button_set_unet_secondary.click(
fn=lambda value, _: run_settings_single(value, key="sd_unet_secondary"),
_js="consumeDesiredUNetName",
inputs=[shared.settings_components["sd_unet_secondary"], dummy_component],
outputs=[shared.settings_components["sd_unet_secondary"], text_settings],
)
def reference_submit(model):
if '@' not in model: # diffusers
modules/video_models/models_def.py
+1 -1
View File
@@ -37,7 +37,7 @@ def getpipe(package, name, _default=None):
cls = getattr(package, name, _default)
return cls
except Exception as e:
log.error(f'Video load: error loading class "{name}" from package "{package.__name__}": {e}')
log.error(f'Load video: error loading class "{name}" from package "{package.__name__}": {e}')
return None
modules/video_models/video_load.py
+11 -8
View File
@@ -22,7 +22,7 @@ loaded_model = None
def load_custom(model_name: str):
log.debug(f'Video load: module=pipe repo="{model_name}" cls=Custom')
log.debug(f'Load video: module=pipe repo="{model_name}" cls=Custom')
if 'veo-3.1' in model_name:
from modules.video_models.google_veo import load_veo
pipe = load_veo(model_name)
@@ -48,7 +48,7 @@ def load_model(selected: models_def.Model):
# shared.sd_model auto-reloads the default checkpoint when model_data.sd_model is None,
# which silently swaps the pipe class behind the name-based cache. Pipe-class mismatch
# is the reliable signal that the cached name no longer maps to the cached object.
log.warning(f'Video load: cached model="{selected.name}" cls={type(shared.sd_model).__name__} mismatch forcing reload')
log.warning(f'Load video: cached model="{selected.name}" cls={type(shared.sd_model).__name__} mismatch forcing reload')
loaded_model = None
if loaded_model == selected.name:
return ''
@@ -104,7 +104,7 @@ def load_model(selected: models_def.Model):
selected.te_folder = 'text_encoder'
selected.te_revision = None
log.debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("transformers")}')
log.debug(f'Load video: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("transformers")}')
kwargs["text_encoder"] = selected.te_cls.from_pretrained(
pretrained_model_name_or_path=selected.te or selected.repo,
subfolder=selected.te_folder,
@@ -125,7 +125,7 @@ def load_model(selected: models_def.Model):
if dit_folder is not None and dit_folder not in kwargs:
# get a new quant arg on every loop to prevent the quant config classes getting entangled
load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("diffusers")}')
log.debug(f'Load video: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)} loader={_loader("diffusers")}')
kwargs[dit_folder] = selected.dit_cls.from_pretrained(
pretrained_model_name_or_path=selected.dit or selected.repo,
subfolder=dit_folder,
@@ -136,7 +136,7 @@ def load_model(selected: models_def.Model):
**offline_args,
)
else:
log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} loader={_loader("diffusers")} skip')
log.debug(f'Load video: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} loader={_loader("diffusers")} skip')
if selected.dit_folder is None:
selected.dit_folder = ['transformer']
@@ -154,7 +154,10 @@ def load_model(selected: models_def.Model):
if selected.repo_cls is None:
shared.sd_model = load_custom(selected.repo)
else:
log.debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}')
log.debug(f'Load video: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}')
print('HERE1')
sd_models.hf_prefetch_configs(selected.repo, {}, 'video')
print('HERE2')
shared.sd_model = selected.repo_cls.from_pretrained(
pretrained_model_name_or_path=selected.repo,
revision=selected.repo_revision,
@@ -168,7 +171,7 @@ def load_model(selected: models_def.Model):
errors.display(e, 'video')
if shared.sd_model is None:
msg = f'Video load: model="{selected.name}" failed'
msg = f'Load video: model="{selected.name}" failed'
log.error(msg)
return msg
@@ -211,7 +214,7 @@ def load_model(selected: models_def.Model):
sd_models.set_diffuser_offload(shared.sd_model)
loaded_model = selected.name
msg = f'Video load: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}'
msg = f'Load video: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}'
log.info(msg)
log.debug(f'Video hijacks: decode={decode} text={text} image={image} slicing={slicing} tiling={tiling} framewise={framewise}')
shared.state.end(jobid)
modules/video_models/video_overrides.py
+2 -2
View File
@@ -36,10 +36,10 @@ def load_override(selected: Model, **load_args):
from diffusers.pipelines.ltx2 import LTX2TextConnectors
ltx2_connectors_cls = LTX2TextConnectors
except ImportError as e:
log.warning(f'Video load: LTX2TextConnectors unavailable ({e}); dedup of LTX-2.3 connectors disabled')
log.warning(f'Load video: LTX2TextConnectors unavailable ({e}); dedup of LTX-2.3 connectors disabled')
if ('LTXVideo 2.3' in selected.name and shared.opts.te_shared_te and ltx2_connectors_cls is not None):
conn_repo = 'OzzyGT/LTX-2.3-sdnq-dynamic-int4' if 'SDNQ' in selected.name else 'OzzyGT/LTX-2.3'
log.debug(f'Video load: module=connectors repo="{conn_repo}" cls={ltx2_connectors_cls.__name__} shared={shared.opts.te_shared_te}')
log.debug(f'Load video: module=connectors repo="{conn_repo}" cls={ltx2_connectors_cls.__name__} shared={shared.opts.te_shared_te}')
kwargs['connectors'] = ltx2_connectors_cls.from_pretrained(
conn_repo,
subfolder='connectors',
modules/video_models/video_ui.py
+1 -1
View File
@@ -41,7 +41,7 @@ def model_change(engine, model):
def model_load(engine, model):
debug(f'Video load: engine="{engine}" model="{model}"')
debug(f'Load video: engine="{engine}" model="{model}"')
selected = get_selected(engine, model)
yield f'Video model loading: {selected.name}'
if selected:
pipelines/anima/anima_lora.py
+4
View File
@@ -51,6 +51,10 @@ ANIMA_PREFIXES = (
"diffusion_model.",
"lora_te_",
"lora_unet_",
# LyCORIS-standalone save format; resolved by the universal passthrough in
# native_adapter.resolve_group_targets and routed to lora_transformer_ by
# network_prefix_for's default arm.
"lycoris_",
)
pipelines/chroma/chroma_lora.py
+84 -9
View File
@@ -4,12 +4,15 @@ Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
(``lora_force_diffusers`` off and ``chroma`` in ``allow_native``).
Entry points, one per family: :func:`try_load_lora` (plus DoRA),
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`.
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`,
:func:`try_load_ia3`, :func:`try_load_glora`, :func:`try_load_norm`,
:func:`try_load_full`.
Recognized key prefixes: ``diffusion_model.``, ``transformer.``,
``lora_unet_``, plus bare BFL paths (``double_blocks.`` / ``single_blocks.``)
and bare diffusers paths (``transformer_blocks.`` /
``single_transformer_blocks.`` / ``distilled_guidance_layer.``).
``lora_unet_``, ``lycoris_``, plus bare BFL paths (``double_blocks.`` /
``single_blocks.`` / ``img_in.`` / ``txt_in.`` / ``final_layer.`` /
``distilled_guidance_layer.``) and bare diffusers paths
(``transformer_blocks.`` / ``single_transformer_blocks.``).
Chroma LoRAs are trained against the Flux block layout regardless of save
format. The diffusers ``ChromaTransformer2DModel`` exposes split-attention
@@ -26,8 +29,11 @@ Fused weight handling:
Chroma's modulation generator is the central ``distilled_guidance_layer``
approximator (replacing Flux's per-block ``norm1.linear``). The pruned
``ChromaAdaLayerNormZeroPruned`` classes have no ``.linear`` submodule, so
any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. LoRAs
targeting the approximator pass through unchanged.
any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. Approximator
keys pass through verbatim except the MLP leaves, where BFL
``layers.{i}.in_layer`` / ``out_layer`` rename to diffusers ``linear_1`` /
``linear_2``. Non-block embedder and final-layer keys rename via
``CHROMA_EXTRA_MAP``.
"""
from modules.lora import native_adapter
@@ -38,11 +44,17 @@ from modules.lora.native_adapter import ChunkSpec
KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
BARE_FLUX_PREFIXES = ("double_blocks.", "single_blocks.")
# distilled_guidance_layer. is deliberately a bare-BFL prefix, not a
# bare-diffusers one: the resolver renames BFL MLP leaves (in_layer/out_layer)
# and passes diffusers-named leaves (linear_1/linear_2, in_proj, out_proj,
# norms.N) through verbatim, so both namings route correctly.
BARE_FLUX_PREFIXES = (
"double_blocks.", "single_blocks.",
"img_in.", "txt_in.", "final_layer.", "distilled_guidance_layer.",
)
BARE_DIFFUSERS_PREFIXES = (
"transformer_blocks.", "single_transformer_blocks.",
"distilled_guidance_layer.",
)
@@ -61,11 +73,19 @@ LORA_SUFFIXES = native_adapter.LORA_SUFFIXES
LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES
LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES
OFT_SUFFIXES = native_adapter.OFT_SUFFIXES
IA3_SUFFIXES = native_adapter.IA3_SUFFIXES
GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES
NORM_SUFFIXES = native_adapter.NORM_SUFFIXES
FULL_SUFFIXES = native_adapter.FULL_SUFFIXES
LORA_MARKERS = native_adapter.LORA_MARKERS
LOKR_MARKERS = native_adapter.LOKR_MARKERS
LOHA_MARKERS = native_adapter.LOHA_MARKERS
OFT_MARKERS = native_adapter.OFT_MARKERS
IA3_MARKERS = native_adapter.IA3_MARKERS
GLORA_MARKERS = native_adapter.GLORA_MARKERS
NORM_MARKERS = native_adapter.NORM_MARKERS
FULL_MARKERS = native_adapter.FULL_MARKERS
SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE
BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED
@@ -116,6 +136,13 @@ def resolve_targets(prefix_used, base):
def _kohya_to_diffusers(base):
"""For kohya keys like ``double_blocks_0_img_attn_qkv`` or ``single_blocks_5_linear1``."""
extra = CHROMA_EXTRA_KOHYA_MAP.get(base)
if extra is not None:
return [(extra, None)]
if base.startswith("distilled_guidance_layer_layers_"):
for bfl_leaf, dif_leaf in GUIDANCE_LEAF_MAP.items():
if base.endswith("_" + bfl_leaf):
return [(base[:-len(bfl_leaf)] + dif_leaf, None)]
if base.startswith("double_blocks_"):
rest = base[len("double_blocks_"):]
idx, _, suffix = rest.partition("_")
@@ -129,6 +156,14 @@ def _kohya_to_diffusers(base):
def _bfl_to_diffusers(base):
"""For BFL dotted keys like ``double_blocks.0.img_attn.qkv``."""
extra = CHROMA_EXTRA_MAP.get(base)
if extra is not None:
return [(extra, None)]
if base.startswith("distilled_guidance_layer.layers."):
stem, _, leaf = base.rpartition(".")
mapped = GUIDANCE_LEAF_MAP.get(leaf)
if mapped is not None:
return [(f"{stem}.{mapped}", None)]
parts = base.split(".")
if len(parts) < 3:
return [(base, None)]
@@ -142,6 +177,27 @@ def _bfl_to_diffusers(base):
return [(base, None)]
# Non-block BFL targets. Chroma prunes time_in/guidance_in/vector_in and the
# final-layer adaLN modulation, so only the embedders and the final projection
# need renames. None carry a block index, so the kohya form is derivable by
# underscoring the BFL path.
CHROMA_EXTRA_MAP = {
"img_in": "x_embedder",
"txt_in": "context_embedder",
"final_layer.linear": "proj_out",
}
CHROMA_EXTRA_KOHYA_MAP = {k.replace(".", "_"): v for k, v in CHROMA_EXTRA_MAP.items()}
# distilled_guidance_layer MLP leaves: BFL in_layer/out_layer vs diffusers
# PixArtAlphaTextProjection linear_1/linear_2. The other approximator leaves
# (in_proj, out_proj, norms.N) share names on both sides and pass verbatim.
GUIDANCE_LEAF_MAP = {
"in_layer": "linear_1",
"out_layer": "linear_2",
}
# Static non-fused renames (underscore-keyed for dispatch from either kohya or BFL paths).
_DOUBLE_STATIC = {
"img_attn_proj": "attn.to_out.0",
@@ -228,9 +284,28 @@ def try_load_oft(name, network_on_disk, lora_scale):
return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_ia3(name, network_on_disk, lora_scale):
return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_glora(name, network_on_disk, lora_scale):
return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_norm(name, network_on_disk, lora_scale):
return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_full(name, network_on_disk, lora_scale):
return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load(name, network_on_disk, lora_scale):
"""Run every Chroma family loader, merge any that match."""
return native_adapter.try_load_chain(
name, network_on_disk, lora_scale,
family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft),
family_loaders=(
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
),
)
pipelines/ernie/ernie_lora.py
+31 -2
View File
@@ -4,7 +4,9 @@ Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
(``lora_force_diffusers`` off and ``ernieimage`` in ``allow_native``).
Entry points, one per family: :func:`try_load_lora` (plus DoRA),
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`.
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`,
:func:`try_load_ia3`, :func:`try_load_glora`, :func:`try_load_norm`,
:func:`try_load_full`.
Recognized key prefixes: ``diffusion_model.``, ``transformer.``,
``lora_unet_``, plus bare diffusers paths (``layers.``, ``adaLN_modulation.``,
@@ -34,11 +36,19 @@ LORA_SUFFIXES = native_adapter.LORA_SUFFIXES
LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES
LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES
OFT_SUFFIXES = native_adapter.OFT_SUFFIXES
IA3_SUFFIXES = native_adapter.IA3_SUFFIXES
GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES
NORM_SUFFIXES = native_adapter.NORM_SUFFIXES
FULL_SUFFIXES = native_adapter.FULL_SUFFIXES
LORA_MARKERS = native_adapter.LORA_MARKERS
LOKR_MARKERS = native_adapter.LOKR_MARKERS
LOHA_MARKERS = native_adapter.LOHA_MARKERS
OFT_MARKERS = native_adapter.OFT_MARKERS
IA3_MARKERS = native_adapter.IA3_MARKERS
GLORA_MARKERS = native_adapter.GLORA_MARKERS
NORM_MARKERS = native_adapter.NORM_MARKERS
FULL_MARKERS = native_adapter.FULL_MARKERS
SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE
BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED
@@ -106,9 +116,28 @@ def try_load_oft(name, network_on_disk, lora_scale):
return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_ia3(name, network_on_disk, lora_scale):
return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_glora(name, network_on_disk, lora_scale):
return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_norm(name, network_on_disk, lora_scale):
return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_full(name, network_on_disk, lora_scale):
return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load(name, network_on_disk, lora_scale):
"""Run every ERNIE family loader, merge any that match."""
return native_adapter.try_load_chain(
name, network_on_disk, lora_scale,
family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft),
family_loaders=(
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
),
)
pipelines/flux/flux2_lora.py
+36 -16
View File
@@ -15,7 +15,8 @@ produced by ``Flux2Transformer2DModel.save_lora_adapter()``). Diffusers-PEFT
``.lora_A.default.weight``) is stripped to match the standard suffix table.
BFL/kohya keys are mapped to diffusers paths via ``F2_SINGLE_MAP`` /
``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks emits three
``F2_DOUBLE_MAP`` / ``F2_QKV_MAP`` / ``F2_EXTRA_MAP`` (non-block targets:
embedders, modulation, final layer). Fused QKV in double_blocks emits three
Q/K/V targets each carrying a :class:`modules.lora.native_adapter.ChunkSpec`
that the generic loaders use to chunk the up-weight or instantiate the
appropriate ``NetworkModule*Chunk`` variant.
@@ -55,11 +56,11 @@ from modules.lora.native_adapter import ChunkSpec
# === Arch-specific prefix configuration ===
KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT + ("lycoris_",)
KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
BARE_FLUX_PREFIXES = (
"single_blocks.", "double_blocks.", "img_in.", "txt_in.",
"final_layer.", "time_in.", "single_stream_modulation.",
"final_layer.", "time_in.", "guidance_in.", "single_stream_modulation.",
"double_stream_modulation_",
)
@@ -92,6 +93,28 @@ F2_QKV_MAP = {
"txt_attn.qkv": ("attn", ["add_q_proj", "add_k_proj", "add_v_proj"]),
}
# Non-block BFL targets: embedders, timestep/guidance MLPs, modulation and the
# final layer. Exact-match on the full base. guidance_in resolves only on models
# built with guidance_embeds (the module is absent otherwise and the group is
# counted as unmapped).
F2_EXTRA_MAP = {
"img_in": "x_embedder",
"txt_in": "context_embedder",
"time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
"time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
"guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
"guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
"final_layer.linear": "proj_out",
"final_layer.adaLN_modulation.1": "norm_out.linear",
"single_stream_modulation.lin": "single_stream_modulation.linear",
"double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
"double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
}
# Kohya underscores the full BFL path; non-block paths carry no index so the
# kohya form is derivable.
F2_EXTRA_KOHYA_MAP = {k.replace(".", "_"): v for k, v in F2_EXTRA_MAP.items()}
# Kohya underscore suffix -> BFL dot suffix. Used to convert kohya key fragments
# to look up F2_DOUBLE_MAP / F2_QKV_MAP.
KOHYA_SUFFIX_MAP = {
@@ -160,31 +183,25 @@ def resolve_targets(prefix_used, base):
"""Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group key.
For ``lora_unet_`` prefix, applies ``KOHYA_SUFFIX_MAP`` then ``F2_*_MAP``.
For BFL / bare-BFL, applies ``F2_*_MAP`` directly. ``lycoris_`` is an
already-underscored diffusers path, returned verbatim. Unrecognized
prefixes return an empty list.
For BFL / bare-BFL, applies ``F2_*_MAP`` directly. Unrecognized prefixes
return an empty list.
Universal passthrough prefixes are handled upstream by
:func:`native_adapter.resolve_group_targets`.
Universal passthrough prefixes (including ``lycoris_``) are handled
upstream by :func:`native_adapter.resolve_group_targets`.
"""
if prefix_used == "lora_unet_":
return _kohya_to_diffusers_targets(base)
if prefix_used in (None, "diffusion_model."):
return _bfl_to_diffusers_targets(base)
if prefix_used == "lycoris_":
# base is an already-underscored diffusers path (e.g.
# 'transformer_blocks_0_attn_add_k_proj'). The caller's network_key
# construction does base.replace('.', '_'); for already-underscored
# paths that's a no-op, so the network_key matches the entry stamped
# by lora_convert.assign_network_names_to_compvis_modules
# (e.g. 'lora_transformer_transformer_blocks_0_attn_add_k_proj').
return [(base, None)]
return []
def _kohya_to_diffusers_targets(stripped):
"""For kohya keys like ``double_blocks_0_img_attn_proj`` or ``single_blocks_5_linear1``."""
targets: list[tuple[str, ChunkSpec | None]] = []
extra = F2_EXTRA_KOHYA_MAP.get(stripped)
if extra is not None:
return [(extra, None)]
if stripped.startswith("single_blocks_"):
rest = stripped[len("single_blocks_"):]
idx, _, suffix = rest.partition("_")
@@ -211,6 +228,9 @@ def _kohya_to_diffusers_targets(stripped):
def _bfl_to_diffusers_targets(base):
"""For BFL keys like ``double_blocks.0.img_attn.proj`` or ``single_blocks.5.linear1``."""
targets: list[tuple[str, ChunkSpec | None]] = []
extra = F2_EXTRA_MAP.get(base)
if extra is not None:
return [(extra, None)]
parts = base.split(".")
if len(parts) < 3:
return targets
pipelines/generic_transformer.py
+37 -14
View File
@@ -21,10 +21,16 @@ def load_transformer(
modules_dtype_dict=None,
use_safetensors=True,
native_spec=None,
override_slot='primary',
**kwargs):
"""Load a DiT transformer from the base repo, or from a user-selected
single file when the UNET dropdown (``shared.opts.sd_unet``) is set.
single file when the slot's UNET override dropdown is set.
``override_slot`` selects which dropdown this call consumes: ``'primary'``
reads ``shared.opts.sd_unet``, ``'secondary'`` reads
``shared.opts.sd_unet_secondary`` (dual-transformer arches give each
transformer its own slot).
With ``native_spec`` set and a .safetensors override selected, dispatches
to :func:`pipelines.native_transformer.load`. Without a spec, a single-file
@@ -69,14 +75,23 @@ def load_transformer(
)
local_file = None
override_name = None
fallback = True
from modules import sd_unet
if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
if shared.opts.sd_unet not in list(sd_unet.unet_dict):
log.error(f'Load module: type=transformer file="{shared.opts.sd_unet}" not found')
elif os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]):
local_file = sd_unet.unet_dict[shared.opts.sd_unet]
if override_slot == 'primary':
override_opt, tracker_attr = 'sd_unet', 'loaded_unet'
elif override_slot == 'secondary':
override_opt, tracker_attr = 'sd_unet_secondary', 'loaded_unet_secondary'
else:
raise ValueError(f'load_transformer: unknown override_slot={override_slot}')
selected = getattr(shared.opts, override_opt, None)
if selected is not None and selected != 'Default':
if selected not in list(sd_unet.unet_dict):
log.error(f'Load module: type=transformer slot={override_slot} file="{selected}" not found')
elif os.path.exists(sd_unet.unet_dict[selected]):
local_file = sd_unet.unet_dict[selected]
override_name = selected
if repo_id.startswith(shared.opts.ckpt_dir) and os.path.exists(repo_id):
log.error(f'Load model: transformer="{repo_id}" is incorrectly placed in the checkpoints folder')
@@ -118,8 +133,8 @@ def load_transformer(
log.warning(f'Load model: transformer="{local_file}" override incompatible with cls={cls_name.__name__} ({e})')
if fallback:
log.warning(f'Load model: transformer="{local_file}" ignoring override and loading base transformer')
shared.opts.data['sd_unet'] = 'Default'
sd_unet.loaded_unet = None
shared.opts.data[override_opt] = 'Default'
setattr(sd_unet, tracker_attr, None)
transformer = load_from_repo()
# 3. load safetensors with diffusers loader
@@ -142,6 +157,11 @@ def load_transformer(
else:
transformer = load_from_repo()
# mark the dropdown selection as loaded so the slot's onchange callback
# does not force a redundant full reload for an already-consumed override
if transformer is not None and override_name is not None and getattr(shared.opts, override_opt, None) == override_name:
setattr(sd_unet, tracker_attr, override_name)
sd_models.allow_post_quant = False # we already handled it
if shared.opts.diffusers_offload_mode != 'none' and transformer is not None:
sd_models.move_model(transformer, devices.cpu)
@@ -166,12 +186,15 @@ def load_transformer(
log.debug(f'Load model: transformer="{repo_id}" quant="{quant_type}" size={module_size:.3f} params={param_num:.3f} memory={module_memory}')
try:
actual_dtype = transformer.dtype
if isinstance(actual_dtype, torch.dtype) and isinstance(dtype, torch.dtype) and actual_dtype != dtype:
force = shared.opts.force_dtype
log.warning(f'Load model: transformer="{repo_id}" dtype desired={dtype} actual={actual_dtype} force={force}')
if force:
transformer = transformer.to(dtype)
# quantized models legitimately report the storage dtype (e.g. fp8 comfy_quant
# adopted via SDNQ); the compute dtype lives in the dequantizers, not the params
if getattr(transformer, 'quantization_config', None) is None:
actual_dtype = transformer.dtype
if isinstance(actual_dtype, torch.dtype) and isinstance(dtype, torch.dtype) and actual_dtype != dtype:
force = shared.opts.force_dtype
log.warning(f'Load model: transformer="{repo_id}" dtype desired={dtype} actual={actual_dtype} force={force}')
if force:
transformer = transformer.to(dtype)
except Exception:
pass
pipelines/ideogram/__init__.py
+62
View File
@@ -0,0 +1,62 @@
"""Ideogram 4 native loader spec.
Kept import-light: only the converter and spec live here so the native
loader can import them without pulling the pipeline or the qwen patch.
"""
import re
import torch
import diffusers
from pipelines.native_transformer import TransformerSpec
QKV_TARGETS = ("to_q", "to_k", "to_v")
QKV_RE = re.compile(r"^(layers\.\d+\.attention)\.qkv\.(.+)$")
OUT_RE = re.compile(r"^(layers\.\d+\.attention)\.o\.(.+)$")
def convert_ideogram4_transformer_checkpoint(state_dict: dict, **kwargs) -> dict: # pylint: disable=unused-argument
"""Fused community layout to ``Ideogram4Transformer2DModel`` layout.
``layers.N.attention.qkv.<suffix>`` splits into ``to_q/to_k/to_v.<suffix>``:
row-stacked tensors (weight, row-wise scale, bias) are dim-0 sliced into
thirds, non-row-stacked sidecars (markers, scalar scales) copied verbatim
to all three. ``attention.o.<suffix>`` renames to ``attention.to_out.0.<suffix>``,
everything else passes through. Returns a new dict; the input (which may
be the cached state dict) is not mutated.
"""
converted: dict = {}
for key, value in state_dict.items():
fused_match = QKV_RE.match(key)
if fused_match is not None:
prefix, suffix = fused_match.group(1), fused_match.group(2)
fused_weight = state_dict.get(f"{prefix}.qkv.weight")
fused_rows = fused_weight.shape[0] if torch.is_tensor(fused_weight) and fused_weight.ndim >= 1 else None
row_stacked = (
suffix != "comfy_quant"
and torch.is_tensor(value) and value.ndim >= 1
and fused_rows is not None and value.shape[0] == fused_rows and fused_rows % 3 == 0
)
if row_stacked:
third = fused_rows // 3
for i, target in enumerate(QKV_TARGETS):
converted[f"{prefix}.{target}.{suffix}"] = value[i * third:(i + 1) * third]
else:
for target in QKV_TARGETS:
converted[f"{prefix}.{target}.{suffix}"] = value
continue
out_match = OUT_RE.match(key)
if out_match is not None:
converted[f"{out_match.group(1)}.to_out.0.{out_match.group(2)}"] = value
continue
converted[key] = value
return converted
IDEOGRAM4_SPEC = TransformerSpec(
cls=diffusers.Ideogram4Transformer2DModel,
converter=convert_ideogram4_transformer_checkpoint,
converter_handles_quant=True,
)
pipelines/ideogram/ideogram4.py
+21 -5
View File
@@ -161,6 +161,20 @@ def _expand_tensor_to_effective_batch(
return torch.repeat_interleave(tensor, repeats=repeat_by, dim=0, output_size=tensor.shape[0] * repeat_by)
def transformer_compute_dtype(module: torch.nn.Module) -> torch.dtype:
"""Activation dtype of a possibly-quantized transformer: sub-16-bit
floating params (e.g. fp8) surface through `module.dtype` but are a
storage format; the dequantizer's result dtype is the compute dtype."""
dtype = module.dtype
if not dtype.is_floating_point or torch.finfo(dtype).bits >= 16:
return dtype
for m in module.modules():
dequantizer = getattr(m, 'sdnq_dequantizer', None)
if dequantizer is not None:
return dequantizer.result_dtype
return dtype
class Ideogram4Pipeline(DiffusionPipeline):
r"""
Text-to-image pipeline for Ideogram4.
@@ -675,8 +689,10 @@ class Ideogram4Pipeline(DiffusionPipeline):
# The transformers run in their loaded compute dtype; cast the (otherwise float32) text features to match.
# `latents` stay float32 for scheduler precision and are cast per-step at the transformer call below.
llm_features = llm_features.to(self.transformer.dtype)
neg_llm_features = neg_llm_features.to(self.unconditional_transformer.dtype if self.unconditional_transformer else self.transformer.dtype)
cond_dtype = transformer_compute_dtype(self.transformer)
uncond_dtype = transformer_compute_dtype(self.unconditional_transformer) if self.unconditional_transformer is not None else cond_dtype
llm_features = llm_features.to(cond_dtype)
neg_llm_features = neg_llm_features.to(uncond_dtype)
# 8. Denoising loop. The scheduler stores `num_train_timesteps`-scaled timesteps; convert back to model time.
num_train_timesteps = self.scheduler.config.num_train_timesteps # pylint: disable=no-member
@@ -687,10 +703,10 @@ class Ideogram4Pipeline(DiffusionPipeline):
# Map sigma-domain timestep to model time `t` in [0, 1] (0 = noise, 1 = clean data).
t_model = 1.0 - (t.float() / num_train_timesteps)
t_model = t_model.expand(batch_size * num_images_per_prompt).to(self.transformer.dtype)
t_model = t_model.expand(batch_size * num_images_per_prompt).to(cond_dtype)
# Conditional pass operates on the full packed sequence.
pos_z = torch.cat([text_z_padding, latents], dim=1).to(self.transformer.dtype)
pos_z = torch.cat([text_z_padding, latents], dim=1).to(cond_dtype)
pos_out = self.transformer(
hidden_states=pos_z,
timestep=t_model,
@@ -710,7 +726,7 @@ class Ideogram4Pipeline(DiffusionPipeline):
if gw[i] > 1.0:
uncond_transformer = self.unconditional_transformer if self.unconditional_transformer is not None else self.transformer
neg_v = uncond_transformer(
hidden_states=latents.to(uncond_transformer.dtype),
hidden_states=latents.to(uncond_dtype),
timestep=t_model,
encoder_hidden_states=neg_llm_features,
position_ids=neg_position_ids,
pipelines/krea2/krea2_lora.py
+20 -1
View File
@@ -63,9 +63,28 @@ def try_load_oft(name, network_on_disk, lora_scale):
return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_ia3(name, network_on_disk, lora_scale):
return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_glora(name, network_on_disk, lora_scale):
return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_norm(name, network_on_disk, lora_scale):
return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_full(name, network_on_disk, lora_scale):
return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load(name, network_on_disk, lora_scale):
"""Run every Krea 2 family loader, merge any that match."""
return native_adapter.try_load_chain(
name, network_on_disk, lora_scale,
family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft),
family_loaders=(
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
),
)
pipelines/model_anima.py
+1 -2
View File
@@ -120,8 +120,7 @@ def load_anima(checkpoint_info, diffusers_load_config=None):
repo_id,
cls_name=transformers.Qwen3Model,
load_config=diffusers_load_config,
subfolder="text_encoder",
allow_shared=False
subfolder="text_encoder"
)
if llm_adapter is None:
pipelines/model_ideogram4.py
+9 -2
View File
@@ -1,3 +1,5 @@
import dataclasses
import diffusers
import transformers
from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
@@ -41,11 +43,16 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None):
if repo_id is None or repo_id.lower() == 'none':
return None
from pipelines.ideogram import IDEOGRAM4_SPEC
transformer_cls = diffusers.Ideogram4Transformer2DModel
transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config)
transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config, native_spec=IDEOGRAM4_SPEC)
if shared.opts.model_ideogram4_enable_cg:
unconditional_transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config)
# the spec subfolder swap makes an unconditional override fetch its config from the matching subfolder
unconditional_transformer = generic.load_transformer(
repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config,
native_spec=dataclasses.replace(IDEOGRAM4_SPEC, subfolder="unconditional_transformer"), override_slot='secondary',
)
else:
unconditional_transformer = None
text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLModel, load_config=diffusers_load_config)
pipelines/model_joy.py
+13 -5
View File
@@ -5,7 +5,7 @@ from modules.logger import log
from pipelines import generic
def load_joy(checkpoint_info, diffusers_load_config=None):
def load_joyedit(checkpoint_info, diffusers_load_config=None):
if diffusers_load_config is None:
diffusers_load_config = {}
repo_id = sd_models.path_to_repo(checkpoint_info)
@@ -15,9 +15,13 @@ def load_joy(checkpoint_info, diffusers_load_config=None):
log.debug(f'Load model: type=JoyImageEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
from pipelines.joy import JOY_SPEC
if 'plus' in repo_id.lower():
dit_cls = diffusers.JoyImageEditPlusTransformer3DModel
else:
dit_cls = diffusers.JoyImageEditTransformer3DModel
transformer = generic.load_transformer(
repo_id,
cls_name=diffusers.JoyImageEditTransformer3DModel,
cls_name=dit_cls,
load_config=diffusers_load_config,
native_spec=JOY_SPEC,
)
@@ -29,7 +33,11 @@ def load_joy(checkpoint_info, diffusers_load_config=None):
if repo_id is None or repo_id.lower() == 'none':
return None
pipe = diffusers.JoyImageEditPipeline.from_pretrained(
if 'plus' in repo_id.lower():
model_cls = diffusers.JoyImageEditPlusPipeline
else:
model_cls = diffusers.JoyImageEditPipeline
pipe = model_cls.from_pretrained(
repo_id,
cache_dir=shared.opts.diffusers_dir,
transformer=transformer,
@@ -40,8 +48,8 @@ def load_joy(checkpoint_info, diffusers_load_config=None):
'output_type': 'np',
}
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = diffusers.JoyImageEditPipeline
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = diffusers.JoyImageEditPipeline
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = model_cls
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = model_cls
generic.load_vae_override(pipe, diffusers_load_config)
pipelines/model_wanai.py
+1 -1
View File
@@ -42,7 +42,7 @@ def load_wan(checkpoint_info, diffusers_load_config=None):
boundary_ratio = 1000.0
elif shared.opts.model_wan_stage == 'combined' or shared.opts.model_wan_stage == 'both':
transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer')
transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2')
transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2', override_slot='secondary')
# load with the checkpoint's boundary; the slider override is applied at runtime in set_pipeline_args
boundary_ratio = None
else:
pipelines/native_transformer.py
+548 -20
View File
@@ -19,15 +19,19 @@ from diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table.
Algorithm:
1. Read the safetensors state dict (.gguf and .pth are rejected up front).
2. Detect and strip one of the spec's known prefixes (raises on mixed prefixes).
3. Check forbidden markers (catches structural mismatches like Cosmos 1.0 keys
2. Drop keys belonging to known companion component families (all-in-one
exports bundle the text encoder and VAE under prefixes like
``cond_stage_model.`` / ``first_stage_model.``; sdnext sources those
components elsewhere).
3. Detect and strip one of the spec's known prefixes (raises on mixed prefixes).
4. Check forbidden markers (catches structural mismatches like Cosmos 1.0 keys
in a Cosmos 2.0 loader).
4. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``).
5. Run the spec's converter if present (else pass through unchanged).
6. Fetch ``<subfolder>/config.json`` from the base repo, instantiate via
5. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``).
6. Run the spec's converter if present (else pass through unchanged).
7. Fetch ``<subfolder>/config.json`` from the base repo, instantiate via
``cls.from_config``, ``load_state_dict(strict=False)``, validate, dtype-cast,
quantize, and offload-place.
7. Repeat the build for each populated sibling (no converter, no quant by
8. Repeat the build for each populated sibling (no converter, no quant by
default; sibling weights are read raw from the bundled file).
Returns ``(transformer, sibling_components_dict)``. The dict is empty for
@@ -37,6 +41,7 @@ loading the adapter from the base repo.
"""
import os
import json
import time
from dataclasses import dataclass, field
from typing import Callable
@@ -58,6 +63,20 @@ DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = (
"pos_embedder.",
"learnable_pos_embed.",
)
DEFAULT_IGNORED_PREFIXES: tuple[str, ...] = (
"cond_stage_model.",
"conditioner.",
"first_stage_model.",
"text_encoders.",
"vae.",
)
COMFY_QUANT_MARKER = ".comfy_quant"
COMFY_QUANT_FORMATS: dict[str, str] = { # comfy_quant format string -> SDNQ weights_dtype
"int8_tensorwise": "int8",
"float8_e4m3fn": "float8_e4m3fn",
"nvfp4": "float4_e2m1fn",
}
NVFP4_GROUP_SIZE = 16 # fixed by the format definition; markers restate it
class OverrideArchMismatch(Exception):
@@ -106,12 +125,23 @@ class TransformerSpec:
``acceptable_missing`` (buffer-only keys left at their init), these are
zero-filled on load so the branch stays a no-op, matching a base model
that ships the branch dormant (output projection all-zeros).
``ignored_prefixes`` names companion component families (text encoder,
VAE) that all-in-one exports bundle alongside the transformer; their keys
are dropped before prefix detection rather than treated as a malformed
file.
``converter_handles_quant`` runs the converter before comfy_quant
detection; such converters must translate marker/scale sidecar keys along
with the weights. Float-oriented converters keep the default.
"""
cls: type
subfolder: str = "transformer"
prefixes: tuple[str, ...] = DEFAULT_PREFIXES
ignored_prefixes: tuple[str, ...] = DEFAULT_IGNORED_PREFIXES
converter: Callable[[dict], dict] | None = None
converter_handles_quant: bool = False
siblings: dict[str, SiblingSpec] = field(default_factory=dict)
acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING
zero_init_missing: tuple[str, ...] = ()
@@ -246,7 +276,12 @@ def load(
quant_type = model_quant.get_quant_type(quant_args)
state_dict = sd_models.read_state_dict(local_file, what="transformer")
metadata_layers = read_quantization_metadata(local_file)
state_dict = drop_companion_keys(state_dict, spec.ignored_prefixes, spec.cls.__name__)
state_dict, detected_prefix = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__)
if metadata_layers and detected_prefix:
# header metadata names mirror the file's tensor naming, so they carry the same prefix
metadata_layers = {name[len(detected_prefix):] if name.startswith(detected_prefix) else name: meta for name, meta in metadata_layers.items()}
check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file)
transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings)
del state_dict
@@ -274,6 +309,8 @@ def load(
dtype=effective_dtype,
modules_to_not_convert=modules_to_not_convert,
modules_dtype_dict=modules_dtype_dict,
converter_handles_quant=spec.converter_handles_quant,
metadata_layers=metadata_layers,
**kwargs,
)
del transformer_sd
@@ -309,6 +346,40 @@ def load(
return transformer, loaded_siblings
def drop_companion_keys(state_dict: dict, ignored_prefixes: tuple[str, ...], type_name: str) -> dict:
"""Remove keys belonging to known non-transformer component families.
All-in-one exports bundle the text encoder and VAE alongside the
transformer under LDM/ComfyUI-style family prefixes (``cond_stage_model.``,
``first_stage_model.``, ``text_encoders.``, ``vae.``). Only the transformer
(plus any spec-declared inline siblings) is wanted here; sdnext sources the
other components from the base repo or their own override dropdowns.
Dropped families are logged so the skip is visible in the load log.
Raises ValueError when nothing remains after filtering, meaning the
selected file holds no transformer at all.
"""
if not ignored_prefixes:
return state_dict
dropped: dict[str, int] = {}
kept: dict = {}
for key, value in state_dict.items():
prefix = next((p for p in ignored_prefixes if key.startswith(p)), None)
if prefix is None:
kept[key] = value
else:
dropped[prefix] = dropped.get(prefix, 0) + 1
if not dropped:
return state_dict
counts = " ".join(f"{prefix.rstrip('.')}={count}" for prefix, count in dropped.items())
if not kept:
raise ValueError(
f"Load model: type={type_name} native_transformer has no transformer keys ({counts})"
)
log.info(f"Load model: type={type_name} native_transformer skipping bundled components: {counts}")
return kept
def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> tuple[dict, str]:
"""Detect and uniformly strip the most common known prefix from every key.
@@ -363,6 +434,225 @@ def check_forbidden_markers(
)
def read_quantization_metadata(local_file: str) -> dict[str, dict] | None:
"""Read per-layer quantization info from the safetensors header metadata.
Newer comfy_quant checkpoints record layer formats in the header
``__metadata__`` under ``_quantization_metadata`` (a JSON string with a
``layers`` map keyed by tensor naming) instead of per-layer marker
tensors. Returns the layers map, or ``None`` when the header carries no
such entry. Raises :class:`OverrideArchMismatch` when the entry is
present but malformed, so the caller's base-repo fallback engages.
"""
from modules.model_probe import read_safetensors_header
try:
header = read_safetensors_header(local_file)
except Exception as e:
log.debug(f'Load model: file="{local_file}" header metadata unreadable ({e})')
return None
meta = (header.get("__metadata__") or {}).get("_quantization_metadata")
if meta is None:
return None
try:
parsed = json.loads(meta) if isinstance(meta, str) else meta
layers = parsed["layers"]
except Exception as e:
raise OverrideArchMismatch(
f'Load model: file="{local_file}" native_transformer _quantization_metadata '
f"is malformed ({type(e).__name__}: {e})"
) from e
if not isinstance(layers, dict) or not all(isinstance(entry, dict) for entry in layers.values()):
raise OverrideArchMismatch(
f'Load model: file="{local_file}" native_transformer _quantization_metadata '
f"layers map is malformed"
)
return layers
def transcode_quant_metadata(state_dict: dict, metadata_layers: dict[str, dict]) -> tuple[dict, str]:
"""Synthesize per-layer marker tensors from header quantization metadata.
Header metadata and marker tensors describe the same per-layer format
dicts; converging on markers lets detection, converters, and remapping
handle both forms identically. Entries without a matching ``.weight``
(companion components filtered earlier) are skipped; existing markers
are overwritten by their header entry. Returns a new dict plus the
detection source (``header`` or ``both``) for logging.
"""
sd = dict(state_dict)
had_markers = any(key.endswith(COMFY_QUANT_MARKER) for key in sd)
count = 0
for name, meta in metadata_layers.items():
if f"{name}.weight" not in sd:
continue
sd[f"{name}{COMFY_QUANT_MARKER}"] = torch.tensor(list(json.dumps(meta).encode("utf-8")), dtype=torch.uint8)
count += 1
if count == 0:
return state_dict, "markers"
return sd, "both" if had_markers else "header"
def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict], str] | None:
"""Detect ``comfy_quant`` pre-quantized layers in a state dict.
ComfyUI-format quantized checkpoints mark each quantized layer with a
``<name>.comfy_quant`` uint8 tensor whose bytes are a JSON object naming
the storage format, alongside a ``<name>.weight_scale`` tensor. Returns
a mapping of marked module names to their parsed marker JSON plus the
file's format string, or ``None`` when no markers are present. ConvRot
markers (``convrot``/``convrot_groupsize``) are accepted per layer: the
rotation is the same regular Hadamard SDNQ implements, so flagged layers
map onto ``use_hadamard`` (regular Hadamards only exist for power-of-4
group sizes). Raises :class:`OverrideArchMismatch` for
malformed markers, unsupported formats or group sizes, or a file mixing
formats across layers, so the caller's base-repo fallback engages.
"""
marker_keys = [key for key in state_dict if key.endswith(COMFY_QUANT_MARKER)]
if not marker_keys:
return None
marked: dict[str, dict] = {}
formats: set[str] = set()
for key in marker_keys:
name = key[: -len(COMFY_QUANT_MARKER)]
try:
meta = json.loads(state_dict[key].cpu().numpy().tobytes())
fmt = meta["format"]
except Exception as e:
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant marker "
f"for {name!r} is malformed ({type(e).__name__}: {e})"
) from e
if meta.get("convrot"):
if fmt == "nvfp4":
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant layer "
f"{name!r} combines nvfp4 with convrot, which the format does not define"
)
group_size = int(meta.get("convrot_groupsize", 256))
is_pow4 = group_size >= 4 and (group_size & (group_size - 1)) == 0 and (group_size.bit_length() & 1) == 1
if not is_pow4:
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant layer "
f"{name!r} convrot_groupsize={group_size} is not a power of 4"
)
if fmt == "nvfp4" and int(meta.get("group_size", NVFP4_GROUP_SIZE)) != NVFP4_GROUP_SIZE:
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant layer "
f"{name!r} nvfp4 group_size={meta.get('group_size')} is not {NVFP4_GROUP_SIZE}"
)
marked[name] = meta
formats.add(fmt)
unsupported = sorted(formats - set(COMFY_QUANT_FORMATS))
if unsupported:
log.error(
f'Load model: type={type_name} quant=comfy format={",".join(unsupported)} not supported '
f'(supported: {",".join(COMFY_QUANT_FORMATS)})'
)
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant format "
f"{', '.join(unsupported)} not supported"
)
if len(formats) > 1:
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer comfy_quant mixes formats "
f"across layers ({', '.join(sorted(formats))})"
)
return marked, next(iter(formats))
def remap_comfy_quant(state_dict: dict, marked_names: set[str], defer_scales: bool = False) -> dict:
"""Translate comfy_quant tensor naming to SDNQ naming.
Renames ``<name>.weight_scale`` to ``<name>.scale`` (reshaped to 2-D,
``[]`` becomes ``[1, 1]``, since SDNQ transposes scales in place) and
drops the ``<name>.comfy_quant`` markers plus optional
``<name>.input_scale`` activation-calibration sidecars (SDNQ re-derives
activation scales dynamically). With ``defer_scales`` the scale rename is
skipped: block-scaled formats (nvfp4) carry swizzled scale tensors whose
transform needs the layer dimensions, so the prequantized builder handles
them per layer. Returns a new dict; the input (which may be the cached
state dict) is not mutated.
"""
scale_suffix = ".weight_scale"
input_scale_suffix = ".input_scale"
remapped: dict = {}
for key, value in state_dict.items():
if key.endswith(COMFY_QUANT_MARKER) and key[: -len(COMFY_QUANT_MARKER)] in marked_names:
continue
if key.endswith(input_scale_suffix) and key[: -len(input_scale_suffix)] in marked_names:
continue
if not defer_scales and key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names:
remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1)
continue
remapped[key] = value
return remapped
def unswizzle_block_scales(scales: torch.Tensor, rows: int, groups: int) -> torch.Tensor:
"""Invert the cuBLAS 2D block-scaling factor layout back to row-major
``[rows, groups]``, dropping the tile alignment padding.
Block-scaled comfy_quant checkpoints store per-group scales pre-tiled for
the hardware kernels: the ``[roundup(rows, 128), roundup(groups, 4)]``
grid is rearranged into 32x16 tiles as described in
https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
"""
row_blocks = -(rows // -128)
col_blocks = -(groups // -4)
tiles = scales.reshape(-1, 32, 4, 4).transpose(1, 2)
tiles = tiles.reshape(row_blocks, col_blocks, 4, 32, 4).reshape(row_blocks, col_blocks, 128, 4)
return tiles.permute(0, 2, 1, 3).reshape(row_blocks * 128, col_blocks * 4)[:rows, :groups]
def adopt_nvfp4_layer(sd: dict, name: str, linear: torch.nn.Linear, component_name: str, meta: dict) -> None:
"""Rewrite one nvfp4 layer's tensors in place into SDNQ layout.
Slices tile-alignment padding off the packed weight, swaps the nibble
order (the container packs the even element into the high nibble, SDNQ
unpacks low-first), unswizzles the e4m3 block scales, and folds the fp32
global scale into them as ``[out, groups, 1]`` fp32 grouped scales.
"""
out_features, in_features = linear.out_features, linear.in_features
if in_features % NVFP4_GROUP_SIZE != 0:
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked module "
f"{name!r} nvfp4 group size {NVFP4_GROUP_SIZE} does not divide in_features={in_features}"
)
orig_shape = meta.get("orig_shape")
if orig_shape is not None and tuple(orig_shape) != (out_features, in_features):
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked module "
f"{name!r} stores orig_shape={tuple(orig_shape)} but the model expects {(out_features, in_features)}"
)
weight = sd.get(f"{name}.weight")
scales = sd.pop(f"{name}.weight_scale", None)
global_scale = sd.pop(f"{name}.weight_scale_2", None)
if scales is None or global_scale is None:
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked module "
f"{name!r} is missing nvfp4 weight_scale/weight_scale_2 tensors"
)
packed_columns = in_features // 2
if weight is None or weight.ndim != 2 or weight.shape[0] < out_features or weight.shape[1] < packed_columns:
found = tuple(weight.shape) if weight is not None else "missing"
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked module "
f"{name!r} packed weight {found} cannot hold {(out_features, packed_columns)}"
)
groups = in_features // NVFP4_GROUP_SIZE
expected_scales = -(out_features // -128) * 128 * -(groups // -4) * 4
if scales.numel() != expected_scales:
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked module "
f"{name!r} block scales hold {scales.numel()} values, expected {expected_scales}"
)
weight = weight[:out_features, :packed_columns]
weight = torch.bitwise_or(torch.bitwise_left_shift(torch.bitwise_and(weight, 15), 4), torch.bitwise_right_shift(weight, 4))
scales = unswizzle_block_scales(scales, out_features, groups)
sd[f"{name}.weight"] = weight
sd[f"{name}.scale"] = scales.to(torch.float32).mul_(global_scale.to(torch.float32)).unsqueeze(-1)
def partition_siblings(
state_dict: dict,
siblings: dict[str, SiblingSpec],
@@ -500,6 +790,213 @@ def build_component_quantized(
return component
def build_component_prequantized(
*,
component_name: str,
state_dict: dict,
config: dict,
cls: type,
marker_meta: dict[str, dict],
comfy_format: str,
dtype,
acceptable_missing: tuple[str, ...],
zero_init_missing: tuple[str, ...] = (),
**kwargs,
) -> object:
"""Build a component from a comfy_quant pre-quantized state dict, mapping
the marked layers onto SDNQ layers without dequantizing.
The supported formats are subsets of SDNQ's symmetric quantization: same
dequant math (``weight * scale``, no zero point), so tensors are adopted
bit-exact. The 8-bit formats share the storage layout directly (unpacked
``[out, in]``); convrot layers map onto SDNQ's Hadamard support, whose
identical regular Hadamard construction lets the dequantizer undo the
stored rotation; nvfp4 layers keep their packed 4-bit codes (nibble order
swapped once) and land on SDNQ's grouped quantization with the block and
global scales folded into fp32 per-group scales.
The file dictates which layers are quantized, independent of the user's
quantization settings, and floating-point SDNQ params are not cast to the
target dtype (the fp32 scales must survive). Layers are assembled in
canonical dequant layout; ``apply_sdnq_options_to_model`` then applies
the user's quantized-matmul settings, matching
``modules.sdnq.loader.load_sdnq_model``.
"""
import rich.progress as rp
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from diffusers.utils import get_module_from_name
from modules.sdnq.common import dtype_dict, check_torch_compile
from modules.sdnq.kernel_wrappers import is_fp8_compile_supported
from modules.sdnq.quantizer import SDNQConfig, SDNQQuantizer
from modules.sdnq.dequantizer import SDNQDequantizer
from modules.sdnq.layers import get_sdnq_wrapper_class
from modules.sdnq.forward import get_forward_func
from modules.sdnq.loader import apply_sdnq_options_to_model
weights_dtype = COMFY_QUANT_FORMATS[comfy_format]
matmul_dtype = "int8" if dtype_dict[weights_dtype]["is_integer"] else "float8_e4m3fn"
storage_dtype = dtype_dict[weights_dtype]["storage_dtype"]
target_dtype = dtype if dtype is not None else devices.dtype
is_nvfp4 = comfy_format == "nvfp4"
# on hardware where compiled graphs cannot touch e4m3 tensors, adopt fp8 weights
# through the uint8-backed codec: same values (NaN codes decode as +/-480), and
# compiled dequant beats the eager native-fp8 fallback by ~6x
remap_fp8_storage = weights_dtype == "float8_e4m3fn" and check_torch_compile() and not is_fp8_compile_supported
if remap_fp8_storage:
weights_dtype = "float8_e4m3fn_sdnq"
marked_names = set(marker_meta)
sd = remap_comfy_quant(state_dict, marked_names, defer_scales=is_nvfp4)
# Civitai relabels these containers freely; trust the marker only as far
# as the stored tensors actually match it.
for name in marked_names:
weight = sd.get(f"{name}.weight")
if weight is None or weight.dtype != storage_dtype:
found = weight.dtype if weight is not None else "missing"
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant format "
f"{comfy_format} expects {storage_dtype} weights but {name!r} has {found}"
)
if remap_fp8_storage:
sd[f"{name}.weight"] = weight.view(torch.uint8)
# layers flagged full_precision_matrix_mult stay on the dequant path even
# when the user enables quantized matmul (exact-name match on the config list)
full_precision_mm = [f"{name}.weight" for name in sorted(marked_names) if marker_meta[name].get("full_precision_matrix_mult")]
quantization_config = SDNQConfig(
weights_dtype=weights_dtype,
quantized_matmul_dtype=matmul_dtype,
group_size=NVFP4_GROUP_SIZE if is_nvfp4 else -1,
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
add_skip_keys=False,
modules_to_not_convert=[],
modules_to_not_use_matmul=full_precision_mm,
)
quantizer = SDNQQuantizer(quantization_config, pre_quantized=True)
quantizer.torch_dtype = target_dtype
with init_empty_weights(include_buffers=False):
component = cls.from_config(config, **kwargs)
dequant_forward = get_forward_func("Linear", matmul_dtype, False)
for name in sorted(marked_names):
try:
parent, child = get_module_from_name(component, name)
linear = getattr(parent, child)
except (AttributeError, ValueError) as e:
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked "
f"module {name!r} not found in {cls.__name__}"
) from e
if not isinstance(linear, torch.nn.Linear):
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked "
f"module {name!r} is {linear.__class__.__name__}, expected Linear"
)
# convrot weights are stored rotated by SDNQ's regular Hadamard; the dequantizer undoes it
use_hadamard = bool(marker_meta[name].get("convrot"))
hadamard_group_size = int(marker_meta[name].get("convrot_groupsize", 256)) if use_hadamard else 256
if use_hadamard and linear.in_features % hadamard_group_size != 0:
raise OverrideArchMismatch(
f"Load model: transformer=native {component_name} comfy_quant marked "
f"module {name!r} convrot_groupsize={hadamard_group_size} does not divide in_features={linear.in_features}"
)
if is_nvfp4:
adopt_nvfp4_layer(sd, name, linear, component_name, marker_meta[name])
layer_shape = torch.Size((linear.out_features, linear.in_features))
linear.sdnq_dequantizer = SDNQDequantizer(
result_dtype=target_dtype,
result_shape=layer_shape if is_nvfp4 else None,
original_shape=layer_shape,
original_stride=(linear.in_features, 1),
quantized_weight_shape=torch.Size((linear.out_features, linear.in_features // NVFP4_GROUP_SIZE, NVFP4_GROUP_SIZE)) if is_nvfp4 else layer_shape,
weights_dtype=weights_dtype,
quantized_matmul_dtype=matmul_dtype,
hadamard_group_size=hadamard_group_size,
group_size=NVFP4_GROUP_SIZE if is_nvfp4 else -1,
svd_rank=32,
svd_steps=8,
use_quantized_matmul=False,
re_quantize_for_matmul=is_nvfp4,
use_stochastic_rounding=False,
use_hadamard=use_hadamard,
layer_class_name="Linear",
)
wrapped = get_sdnq_wrapper_class(linear, dequant_forward)
wrapped.scale = torch.nn.Parameter(torch.empty((1, 1), dtype=torch.float32, device="meta"), requires_grad=False)
wrapped.zero_point = None
wrapped.svd_up = None
wrapped.svd_down = None
setattr(parent, child, wrapped)
target_device = (
devices.cpu if shared.opts.diffusers_offload_mode != "none"
else devices.device
)
expected_keys = set(component.state_dict().keys())
loaded_keys: set[str] = set()
unexpected: list[str] = []
total = len(sd)
pbar = rp.Progress(
rp.TextColumn(f'[cyan]Load {component_name}:'),
rp.BarColumn(),
rp.MofNCompleteColumn(),
rp.TaskProgressColumn(),
rp.TimeRemainingColumn(),
rp.TimeElapsedColumn(),
rp.TextColumn('[cyan]{task.description}'),
console=console,
)
with pbar:
task = pbar.add_task(total=total, description=cls.__name__)
for name, value in sd.items():
if name in expected_keys:
if quantizer.check_if_quantized_param(component, value, name):
quantizer.create_quantized_param(component, value, name, target_device, dtype=target_dtype)
else:
if torch.is_floating_point(value):
value = value.to(target_dtype)
set_module_tensor_to_device(component, name, target_device, value=value, dtype=target_dtype)
loaded_keys.add(name)
else:
unexpected.append(name)
pbar.update(task, advance=1)
missing = sorted(expected_keys - loaded_keys)
missing = materialize_zero_init(
component, missing, zero_init_missing, device=target_device, dtype=target_dtype
)
validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access
component = apply_sdnq_options_to_model(
component,
dtype=target_dtype,
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
)
return component
def apply_converter(converter: Callable[[dict], dict], state_dict: dict, cls: type, component_name: str) -> dict:
"""Run a spec converter, wrapping any failure as
:class:`OverrideArchMismatch` (with the original chained) so a wrong-arch
file degrades to the base-repo fallback instead of a raw converter crash.
"""
log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}')
try:
return converter(state_dict)
except Exception as e:
raise OverrideArchMismatch(
f"Load model: type={cls.__name__} native_transformer converter "
f"{converter.__name__} rejected the override ({type(e).__name__}: {e}); "
f"file does not look like a {cls.__name__} checkpoint"
) from e
def build_component(
*,
component_name: str,
@@ -514,16 +1011,24 @@ def build_component(
dtype=None,
modules_to_not_convert: list | None = None,
modules_dtype_dict: dict | None = None,
converter_handles_quant: bool = False,
metadata_layers: dict[str, dict] | None = None,
**kwargs,
) -> object:
"""Convert (if needed), instantiate, load weights, dtype-cast, quantize,
and offload-place a single component. Raises on any hard failure.
For the transformer component under SDNQ, the per-tensor pre-mode path
in :func:`build_component_quantized` is used so quantization is applied
in flight (one layer's worth of bf16 in memory at a time). All other
cases (siblings, non-quantized loads, NVIDIAModelOptConfig, layerwise
quant) go through the standard load_state_dict + post-quantize path.
Transformer state dicts carrying ``comfy_quant`` markers dispatch
to :func:`build_component_prequantized` (``quant_args`` are bypassed: the
file is already quantized). ``metadata_layers`` (header-metadata quant
info) is transcoded into markers first, so both container forms share
one path; a ``converter_handles_quant`` converter runs
before that detection, float-oriented converters after it. Under SDNQ the
transformer uses the per-tensor pre-mode path in
:func:`build_component_quantized` so quantization is applied in flight.
All other cases (siblings, non-quantized loads, NVIDIAModelOptConfig,
layerwise quant) go through the standard load_state_dict + post-quantize
path.
``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global
default is used. ``modules_to_not_convert`` and ``modules_dtype_dict``
@@ -532,16 +1037,39 @@ def build_component(
reach ``cls.from_config`` for both construction paths.
"""
try:
quant_source = "markers"
if component_name == "transformer" and metadata_layers:
state_dict, quant_source = transcode_quant_metadata(state_dict, metadata_layers)
if converter is not None and converter_handles_quant and component_name == "transformer":
state_dict = apply_converter(converter, state_dict, cls, component_name)
converter = None # consumed; must not run again on the non-comfy path below
comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None
if comfy_quant is not None:
marker_meta, comfy_format = comfy_quant
convrot_count = sum(1 for meta in marker_meta.values() if meta.get("convrot"))
log.info(
f'Load model: transformer=native {component_name} quant=comfy '
f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} source={quant_source} keys={len(state_dict)} cls={cls.__name__}'
)
component = build_component_prequantized(
component_name=component_name,
state_dict=state_dict,
config=config,
cls=cls,
marker_meta=marker_meta,
comfy_format=comfy_format,
dtype=dtype,
acceptable_missing=acceptable_missing,
zero_init_missing=zero_init_missing,
**kwargs,
)
devices.torch_gc()
return component
if converter is not None:
log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}')
try:
sd = converter(state_dict)
except Exception as e:
raise OverrideArchMismatch(
f"Load model: type={cls.__name__} native_transformer converter "
f"{converter.__name__} rejected the override ({type(e).__name__}: {e}); "
f"file does not look like a {cls.__name__} checkpoint"
) from e
sd = apply_converter(converter, state_dict, cls, component_name)
else:
sd = state_dict
pipelines/z_image/zimage_lora.py
+88 -2
View File
@@ -4,12 +4,22 @@ Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
(``lora_force_diffusers`` off and ``zimage`` in ``allow_native``).
Entry points, one per family: :func:`try_load_lora` (plus DoRA),
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`.
:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`,
:func:`try_load_ia3`, :func:`try_load_glora`, :func:`try_load_norm`,
:func:`try_load_full`.
Recognized key prefixes: ``diffusion_model.``, ``transformer.``,
``lora_unet_``, or bare paths starting with the known block-level prefixes
(``layers.``, ``noise_refiner.``, ``context_refiner.``).
Checkpoint names that differ from the diffusers module tree are rewritten by
:func:`resolve_targets`: qk-norms (``attention.q_norm`` / ``k_norm`` ->
``attention.norm_q`` / ``norm_k``) and the non-block targets in
``ZIMAGE_EXTRA_MAP`` (``x_embedder`` and ``final_layer.*`` live in ModuleDicts
keyed by ``"{patch_size}-{f_patch_size}"``, read from the live model by
:func:`patch_keys`). ``t_embedder.mlp.N`` and ``cap_embedder.N`` already match
and pass through.
Pre-refactor Z-Image attention layouts (fused ``attention.qkv``, bare
``attention.out`` / ``attention.wo``) are rewritten by :func:`resolve_targets`
to the current diffusers ``to_q``/``to_k``/``to_v`` and ``to_out.0``. For
@@ -32,6 +42,42 @@ KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
BARE_DIFFUSERS_PREFIXES = ("layers.", "noise_refiner.", "context_refiner.")
# Checkpoint qk-norm names vs the diffusers attention module names.
ZIMAGE_NORM_ALIASES = {
".attention.q_norm": ".attention.norm_q",
".attention.k_norm": ".attention.norm_k",
}
# Non-block targets. The patch embedder and the final layer live in ModuleDicts keyed by
# "{patch_size}-{f_patch_size}", so the diffusers path carries a key the checkpoint does not:
# {key} is filled from the live model. t_embedder.mlp.N and cap_embedder.N already match and
# pass through verbatim.
ZIMAGE_EXTRA_MAP = {
"x_embedder": "all_x_embedder.{key}",
"final_layer.linear": "all_final_layer.{key}.linear",
"final_layer.adaLN_modulation.1": "all_final_layer.{key}.adaLN_modulation.1",
}
ZIMAGE_EXTRA_KOHYA_MAP = {k.replace(".", "_"): v for k, v in ZIMAGE_EXTRA_MAP.items()}
# Both shipped Z-Image variants (Base, Turbo) build a single entry; read from the live model and
# fall back to it. Tests patch this directly.
PATCH_KEYS_DEFAULT = ["2-1"]
def patch_keys():
"""ModuleDict keys of the live transformer's ``all_x_embedder``."""
try:
from modules import shared
pipe = getattr(shared.sd_model, "pipe", shared.sd_model)
embedder = getattr(getattr(pipe, "transformer", None), "all_x_embedder", None)
keys = list(embedder.keys()) if embedder is not None else []
if keys:
return keys
except Exception:
pass
return PATCH_KEYS_DEFAULT
# === Re-exports for test/back-compat ===
@@ -39,11 +85,19 @@ LORA_SUFFIXES = native_adapter.LORA_SUFFIXES
LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES
LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES
OFT_SUFFIXES = native_adapter.OFT_SUFFIXES
IA3_SUFFIXES = native_adapter.IA3_SUFFIXES
GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES
NORM_SUFFIXES = native_adapter.NORM_SUFFIXES
FULL_SUFFIXES = native_adapter.FULL_SUFFIXES
LORA_MARKERS = native_adapter.LORA_MARKERS
LOKR_MARKERS = native_adapter.LOKR_MARKERS
LOHA_MARKERS = native_adapter.LOHA_MARKERS
OFT_MARKERS = native_adapter.OFT_MARKERS
IA3_MARKERS = native_adapter.IA3_MARKERS
GLORA_MARKERS = native_adapter.GLORA_MARKERS
NORM_MARKERS = native_adapter.NORM_MARKERS
FULL_MARKERS = native_adapter.FULL_MARKERS
SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE
BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED
@@ -96,6 +150,9 @@ def resolve_targets(prefix_used, base):
def _dotted_to_diffusers_targets(base):
"""For BFL / bare-BFL keys like ``layers.0.attention.qkv``."""
extra = ZIMAGE_EXTRA_MAP.get(base)
if extra is not None:
return [(extra.format(key=key), None) for key in patch_keys()]
if base.endswith(".attention.qkv"):
stem = base[:-len(".attention.qkv")]
return [
@@ -103,6 +160,9 @@ def _dotted_to_diffusers_targets(base):
(f"{stem}.attention.to_k", ChunkSpec(idx=1, total=3)),
(f"{stem}.attention.to_v", ChunkSpec(idx=2, total=3)),
]
for alias, renamed in ZIMAGE_NORM_ALIASES.items():
if base.endswith(alias):
return [(base[:-len(alias)] + renamed, None)]
for alias in (".attention.out.0", ".attention.out", ".attention.wo"):
if base.endswith(alias):
stem = base[:-len(alias)]
@@ -112,6 +172,9 @@ def _dotted_to_diffusers_targets(base):
def _underscore_to_diffusers_targets(base):
"""For kohya flat-underscore keys like ``layers_0_attention_qkv``."""
extra = ZIMAGE_EXTRA_KOHYA_MAP.get(base)
if extra is not None:
return [(extra.format(key=key), None) for key in patch_keys()]
if base.endswith("_attention_qkv"):
stem = base[:-len("_attention_qkv")]
return [
@@ -119,6 +182,10 @@ def _underscore_to_diffusers_targets(base):
(f"{stem}_attention_to_k", ChunkSpec(idx=1, total=3)),
(f"{stem}_attention_to_v", ChunkSpec(idx=2, total=3)),
]
for alias, renamed in ZIMAGE_NORM_ALIASES.items():
underscored = alias.replace(".", "_")
if base.endswith(underscored):
return [(base[:-len(underscored)] + renamed.replace(".", "_"), None)]
for alias in ("_attention_out_0", "_attention_out", "_attention_wo"):
if base.endswith(alias):
stem = base[:-len(alias)]
@@ -153,9 +220,28 @@ def try_load_oft(name, network_on_disk, lora_scale):
return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_ia3(name, network_on_disk, lora_scale):
return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_glora(name, network_on_disk, lora_scale):
return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_norm(name, network_on_disk, lora_scale):
return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load_full(name, network_on_disk, lora_scale):
return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS)
def try_load(name, network_on_disk, lora_scale):
"""Run every Z-Image family loader, merge any that match."""
return native_adapter.try_load_chain(
name, network_on_disk, lora_scale,
family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft),
family_loaders=(
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
),
)
pnpm-workspace.yaml
+8
View File
@@ -0,0 +1,8 @@
allowBuilds:
'@google/genai': false
esbuild: true
protobufjs: false
unrs-resolver: false
packages:
- "extensions-builtin/sdnext-modernui"
- "extensions-builtin/sdnext-kanvas"
requirements.txt
+1 -1
View File
@@ -30,7 +30,7 @@ requests==2.34.2
tqdm==4.68.3
accelerate==1.14.0
einops==0.8.2
huggingface_hub==1.22.0
huggingface_hub==1.23.0
hf_xet==1.5.1
numpy==2.1.2
pandas==2.3.1
test/test-chroma-native-adapters.py
+99 -2
View File
@@ -218,10 +218,24 @@ def build_mock_transformer():
transformer = _Holder()
transformer.transformer_blocks = torch.nn.ModuleList([build_double_block() for _ in range(N_DOUBLE)])
transformer.single_transformer_blocks = torch.nn.ModuleList([build_single_block() for _ in range(N_SINGLE)])
# distilled_guidance_layer - Chroma's central modulation approximator
# Minimal stand-in: just one linear submodule the tests can target
# distilled_guidance_layer - Chroma's central modulation approximator.
# Mirrors ChromaApproximator: in_proj / out_proj Linears, PixArt-shaped
# MLP layers (linear_1 / linear_2) and RMSNorms.
transformer.distilled_guidance_layer = _Holder()
transformer.distilled_guidance_layer.in_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.out_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.layers = torch.nn.ModuleList()
transformer.distilled_guidance_layer.norms = torch.nn.ModuleList()
for _ in range(2):
mlp = _Holder()
mlp.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
mlp.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.layers.append(mlp)
transformer.distilled_guidance_layer.norms.append(torch.nn.RMSNorm(HIDDEN))
# Non-block CHROMA_EXTRA_MAP targets.
transformer.x_embedder = torch.nn.Linear(HIDDEN, HIDDEN)
transformer.context_embedder = torch.nn.Linear(HIDDEN, HIDDEN)
transformer.proj_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
return transformer
@@ -630,6 +644,32 @@ def test_resolve_targets_static_renames():
return True
def test_resolve_targets_extra_and_guidance():
"""Non-block extra-map renames and guidance-layer MLP leaf renames, all key forms."""
for bfl_base, diffusers_path in C.CHROMA_EXTRA_MAP.items():
for prefix, base in [
('diffusion_model.', bfl_base),
(None, bfl_base),
('lora_unet_', bfl_base.replace('.', '_')),
]:
targets = C.resolve_targets(prefix, base)
assert targets == [(diffusers_path, None)], f'({prefix}, {base}) -> {targets}'
cases = [
# BFL MLP leaves rename to the PixArt projection names.
(('diffusion_model.', 'distilled_guidance_layer.layers.0.in_layer'), 'distilled_guidance_layer.layers.0.linear_1'),
((None, 'distilled_guidance_layer.layers.1.out_layer'), 'distilled_guidance_layer.layers.1.linear_2'),
(('lora_unet_', 'distilled_guidance_layer_layers_0_in_layer'), 'distilled_guidance_layer_layers_0_linear_1'),
(('lora_unet_', 'distilled_guidance_layer_layers_1_out_layer'), 'distilled_guidance_layer_layers_1_linear_2'),
# Verbatim leaves are untouched in either naming.
(('diffusion_model.', 'distilled_guidance_layer.in_proj'), 'distilled_guidance_layer.in_proj'),
((None, 'distilled_guidance_layer.layers.0.linear_1'), 'distilled_guidance_layer.layers.0.linear_1'),
]
for (prefix, base), expected in cases:
targets = C.resolve_targets(prefix, base)
assert targets == [(expected, None)], f'({prefix}, {base}) -> {targets}'
return True
def test_resolve_targets_onetrainer_passthrough():
"""The ``lora_transformer_`` passthrough lives in the shared resolver.
@@ -805,6 +845,60 @@ def test_lora_distilled_guidance():
return True
def sd_lokr_bfl_extra_modules():
"""BFL LoKR spanning the non-block extra targets and guidance MLP leaves.
Full-matrix factors with the ai-toolkit placeholder alpha, mirroring the
layout of real full-preset checkpoints.
"""
bases = [
'img_in', 'txt_in', 'final_layer.linear',
'distilled_guidance_layer.layers.0.in_layer',
'distilled_guidance_layer.layers.1.out_layer',
]
sd = {}
for base in bases:
sd[f'diffusion_model.{base}.lokr_w1'] = torch.randn(LOKR_W1_DIM, LOKR_W1_DIM)
sd[f'diffusion_model.{base}.lokr_w2'] = torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM)
sd[f'diffusion_model.{base}.alpha'] = torch.tensor(9999220736.0)
return sd
def test_lokr_bfl_extra_and_guidance():
"""Embedder/final-layer renames and guidance MLP leaf renames all bind."""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_extra_modules())
assert net is not None and len(net.modules) == 5, f'got {sorted(net.modules) if net else None}'
expected = {
'lora_transformer_x_embedder',
'lora_transformer_context_embedder',
'lora_transformer_proj_out',
'lora_transformer_distilled_guidance_layer_layers_0_linear_1',
'lora_transformer_distilled_guidance_layer_layers_1_linear_2',
}
assert set(net.modules) == expected, f'got {set(net.modules)}'
# Full-matrix factors: the placeholder alpha must not scale.
for nk, mod in net.modules.items():
assert mod.dim is None and mod.calc_scale() == 1.0, f'{nk}: dim={mod.dim} scale={mod.calc_scale()}'
return True
def test_full_diff_chain():
"""Full-diff extraction loads through the chain; fused qkv diff skips."""
sd = {
'diffusion_model.double_blocks.0.img_attn.proj.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.double_blocks.0.img_attn.proj.diff_b': torch.randn(HIDDEN),
'diffusion_model.double_blocks.0.img_attn.qkv.diff': torch.randn(3 * HIDDEN, HIDDEN),
}
net = _load_via(C.try_load, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules, f'got {set(net.modules)}'
mod = next(iter(net.modules.values()))
updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
assert tuple(updown.shape) == (HIDDEN, HIDDEN) and torch.isfinite(updown).all()
assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,)
return True
def test_lora_dora_threading():
"""dora_scale flows into NetworkModuleLora."""
net = _load_via(C.try_load_lora, sd_lora_with_dora_scale())
@@ -979,6 +1073,7 @@ def run_tests():
log.warning('=== Parsing primitives ===')
for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_resolve_targets_static_renames,
test_resolve_targets_extra_and_guidance,
test_resolve_targets_onetrainer_passthrough]:
run_test(CAT_PARSE, fn)
@@ -1000,6 +1095,8 @@ def run_tests():
test_lokr_bfl_img_attn_proj,
test_lokr_bfl_img_attn_qkv_chunked,
test_lokr_bfl_single_linear1_unequal_chunks,
test_lokr_bfl_extra_and_guidance,
test_full_diff_chain,
test_loha_bfl_img_attn_proj,
test_loha_bfl_img_attn_qkv_chunked,
test_oft_bfl_img_attn_proj,
test/test-ernie-native-adapters.py
+17
View File
@@ -565,6 +565,22 @@ def test_oft_lycoris_no_npe():
return True
def test_full_diff_chain():
"""Full-diff extraction loads through the chain (ERNIE has no fused targets)."""
sd = {
'diffusion_model.layers.0.self_attention.to_q.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.layers.0.self_attention.to_q.diff_b': torch.randn(HIDDEN),
}
net = _load_via(E.try_load, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
assert 'lora_transformer_layers_0_self_attention_to_q' in net.modules, f'got {set(net.modules)}'
mod = next(iter(net.modules.values()))
updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
assert tuple(updown.shape) == (HIDDEN, HIDDEN) and torch.isfinite(updown).all()
assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,)
return True
# ============================================================
# Tests - calc_updown shape sanity
# ============================================================
@@ -635,6 +651,7 @@ def run_tests():
test_lokr_bfl_self_attention,
test_loha_bfl_mlp,
test_oft_lycoris_no_npe,
test_full_diff_chain,
]:
run_test(CAT_LOADER, fn)
test/test-flux2-native-adapters.py
+376 -5
View File
@@ -128,9 +128,10 @@ class _Holder(torch.nn.Module):
def build_mock_transformer():
"""Build a torch.nn.Module mimicking Flux2's diffusers-side module tree.
Mirrors the paths in F2_SINGLE_MAP / F2_DOUBLE_MAP / F2_QKV_MAP plus the
RMSNorm targets inside attention so that try_load_norm has something to
bind. Sizes are scaled-down but proportional so chunking math is realistic.
Mirrors the paths in F2_SINGLE_MAP / F2_DOUBLE_MAP / F2_QKV_MAP /
F2_EXTRA_MAP plus the RMSNorm targets inside attention so that
try_load_norm has something to bind. Sizes are scaled-down but
proportional so chunking math is realistic.
"""
transformer = _Holder()
transformer.transformer_blocks = torch.nn.ModuleList()
@@ -163,6 +164,26 @@ def build_mock_transformer():
sblock.attn.to_qkv_mlp_proj = torch.nn.Linear(HIDDEN, SINGLE_FUSED_OUT, bias=False)
sblock.attn.to_out = torch.nn.Linear(SINGLE_OUT_IN, HIDDEN, bias=False)
transformer.single_transformer_blocks.append(sblock)
# Non-block F2_EXTRA_MAP targets. Modulation/norm_out out-dims follow the
# real modules' expansion ratios (3x/6x per mod_param_sets, 2x for AdaLN).
transformer.x_embedder = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.context_embedder = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.time_guidance_embed = _Holder()
transformer.time_guidance_embed.timestep_embedder = _Holder()
transformer.time_guidance_embed.timestep_embedder.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.time_guidance_embed.timestep_embedder.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.time_guidance_embed.guidance_embedder = _Holder()
transformer.time_guidance_embed.guidance_embedder.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.time_guidance_embed.guidance_embedder.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.norm_out = _Holder()
transformer.norm_out.linear = torch.nn.Linear(HIDDEN, 2 * HIDDEN, bias=False)
transformer.proj_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
transformer.single_stream_modulation = _Holder()
transformer.single_stream_modulation.linear = torch.nn.Linear(HIDDEN, 3 * HIDDEN, bias=False)
transformer.double_stream_modulation_img = _Holder()
transformer.double_stream_modulation_img.linear = torch.nn.Linear(HIDDEN, 6 * HIDDEN, bias=False)
transformer.double_stream_modulation_txt = _Holder()
transformer.double_stream_modulation_txt.linear = torch.nn.Linear(HIDDEN, 6 * HIDDEN, bias=False)
return transformer
@@ -308,6 +329,56 @@ def sd_lora_with_dora_scale():
}
def sd_lora_dora_fused_qkv():
"""Kohya LoRA on fused img_attn.qkv with per-output dora_scale (LyCORIS wd_on_out=True)."""
return {
'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(3 * QKV_OUT, RANK_LORA),
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
'lora_unet_double_blocks_0_img_attn_qkv.dora_scale': torch.rand(3 * QKV_OUT, 1) + 0.5,
}
def sd_lora_diff_b_fused_qkv():
"""Kohya LoRA on fused img_attn.qkv with a per-output diff_b bias delta."""
sd = dict(sd_lora_dora_fused_qkv())
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
sd['lora_unet_double_blocks_0_img_attn_qkv.diff_b'] = torch.randn(3 * QKV_OUT)
return sd
def sd_lora_aitk_magnitude_dora(w_base):
"""ai-toolkit DoRA save: lora_A/B + 1-D magnitude row norms, no alpha key.
Targets img_attn.proj -> attn.to_out.0, a square Linear (HIDDEN == QKV_OUT
in the mock): the case where a raw 1-D vector cannot be disambiguated
between per-output and per-input at apply time.
"""
down = torch.randn(RANK_LORA, QKV_OUT)
up = torch.randn(HIDDEN, RANK_LORA)
magnitude = w_base.reshape(HIDDEN, -1).norm(dim=1) # 1-D (out,)
return {
'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': down,
'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': up,
'diffusion_model.double_blocks.1.img_attn.proj.magnitude': magnitude,
}, down, up, magnitude
def sd_lokr_dora_fused_qkv(per_input=False):
"""Kohya LoKR on fused img_attn.qkv with a dora_scale companion.
``per_input=True`` stores the wd_on_out=False orientation ``(1, in)``,
which has no exact per-chunk split and must be skipped by the loader.
"""
sd = dict(sd_lokr_kohya_qkv())
if per_input:
ds = torch.rand(1, HIDDEN) + 0.5
else:
ds = torch.rand(3 * QKV_OUT, 1) + 0.5
sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale'] = ds
return sd
def sd_lokr_bfl_proj():
"""BFL-format LoKR on a non-fused proj target."""
return {
@@ -361,6 +432,35 @@ def sd_lokr_simpletuner_lycoris_style():
}
def sd_lokr_bfl_extra_modules():
"""BFL-format LoKR spanning every non-block target in F2_EXTRA_MAP.
Mirrors the ai-toolkit full-matrix layout (klein_snofs_v1_1): both factors
stored whole (lokr_w1/lokr_w2, no a/b decomposition) with the ~1e10
placeholder alpha. With no rank-decomposed factor the loader must leave
dim=None so the placeholder never scales the delta.
"""
dims = {
'img_in': (HIDDEN, HIDDEN),
'txt_in': (HIDDEN, HIDDEN),
'time_in.in_layer': (HIDDEN, HIDDEN),
'time_in.out_layer': (HIDDEN, HIDDEN),
'guidance_in.in_layer': (HIDDEN, HIDDEN),
'guidance_in.out_layer': (HIDDEN, HIDDEN),
'final_layer.linear': (HIDDEN, HIDDEN),
'final_layer.adaLN_modulation.1': (2 * HIDDEN, HIDDEN),
'single_stream_modulation.lin': (3 * HIDDEN, HIDDEN),
'double_stream_modulation_img.lin': (6 * HIDDEN, HIDDEN),
'double_stream_modulation_txt.lin': (6 * HIDDEN, HIDDEN),
}
sd = {}
for base, (out_dim, in_dim) in dims.items():
sd[f'diffusion_model.{base}.lokr_w1'] = torch.randn(LOKR_W1_DIM, LOKR_W1_DIM)
sd[f'diffusion_model.{base}.lokr_w2'] = torch.randn(out_dim // LOKR_W1_DIM, in_dim // LOKR_W1_DIM)
sd[f'diffusion_model.{base}.alpha'] = torch.tensor(9999220736.0)
return sd
def sd_loha_bfl_proj():
"""BFL-format LoHA on a non-fused proj target."""
return {
@@ -607,6 +707,21 @@ def test_resolve_targets_qkv_chunking():
return True
def test_resolve_targets_extra_modules():
"""Every non-block F2_EXTRA_MAP target maps 1:1 (no chunk) in all three key forms."""
for bfl_base, diffusers_path in F.F2_EXTRA_MAP.items():
targets = F.resolve_targets('diffusion_model.', bfl_base)
assert targets == [(diffusers_path, None)], f'{bfl_base} -> {targets}'
targets = F.resolve_targets(None, bfl_base)
assert targets == [(diffusers_path, None)], f'bare {bfl_base} -> {targets}'
targets = F.resolve_targets('lora_unet_', bfl_base.replace('.', '_'))
assert targets == [(diffusers_path, None)], f'kohya {bfl_base} -> {targets}'
# guidance_in is a bare BFL prefix in its own right.
got = F.parse_key('guidance_in.in_layer.lora_A.weight', F.LORA_SUFFIXES)
assert got == (None, 'guidance_in.in_layer', 'lora_down.weight'), f'bare guidance_in parse -> {got}'
return True
def test_parse_key_peft_wrapper_unwrap():
"""base_model.model. wrapper is stripped before format detection.
@@ -653,8 +768,9 @@ def test_parse_key_lycoris_prefix():
got = F.parse_key(key, suffixes)
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
# resolve_targets: the underscored path is returned verbatim (no chunk).
targets = F.resolve_targets('lycoris_', 'transformer_blocks_0_attn_add_k_proj')
# Resolution: lycoris_ is a universal passthrough handled upstream of the
# arch resolver, so the underscored path is returned verbatim (no chunk).
targets = F.native_adapter.resolve_group_targets(F.resolve_targets, 'lycoris_', 'transformer_blocks_0_attn_add_k_proj')
assert targets == [('transformer_blocks_0_attn_add_k_proj', None)], f'targets={targets}'
return True
@@ -846,6 +962,185 @@ def test_lora_dora_threading():
return True
def test_lora_dora_fused_qkv_sliced():
"""Per-output dora_scale is sliced with the up-weight chunk and applies cleanly."""
net = _load_via(F.try_load_lora, sd_lora_dora_fused_qkv())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
for nk, mod in net.modules.items():
assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
assert_finite(updown, label=nk)
return True
def test_lokr_dora_fused_qkv_sliced():
"""Per-output dora_scale rides the kron chunk; each module sees its own rows."""
net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
for nk, mod in net.modules.items():
assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}'
assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
assert_finite(updown, label=nk)
return True
def test_dora_per_input_fused_skipped():
"""wd_on_out=False dora_scale on a fused target cannot be split; the group is skipped."""
net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv(per_input=True))
assert net is None, f'expected no modules, got {net.modules if net else None}'
return True
def test_lora_diff_b_fused_qkv_sliced():
"""Per-output diff_b slices with the chunk and flows out as ex_bias."""
net = _load_via(F.try_load_lora, sd_lora_diff_b_fused_qkv())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
for nk, mod in net.modules.items():
assert mod.ex_bias is not None and mod.ex_bias.shape[0] == QKV_OUT, \
f'{nk}: ex_bias shape {tuple(mod.ex_bias.shape) if mod.ex_bias is not None else None}'
_updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
assert ex_bias is not None and ex_bias.shape[0] == QKV_OUT, f'{nk}: applied ex_bias {ex_bias.shape if ex_bias is not None else None}'
return True
def test_legacy_bias_fused_skipped():
"""The legacy weight-shaped bias key has no partition on fused targets; group skipped."""
sd = dict(sd_lora_dora_fused_qkv())
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
sd['lora_unet_double_blocks_0_img_attn_qkv.bias'] = torch.randn(3 * QKV_OUT, HIDDEN)
net = _load_via(F.try_load_lora, sd)
assert net is None, f'expected skip, got {net.modules if net else None}'
return True
def test_lora_aitk_magnitude_dora():
"""ai-toolkit `magnitude` converts onto the dora_scale path, (out, 1)-shaped.
Numeric check on a square layer: with the raw 1-D vector the apply-time
orientation detection would fall to per-input and renormalize columns;
the (out, 1) reshape pins per-output. Reference: (W+D) * m/||W+D||_row - W
at scale 1 (ai-toolkit DoRA saves no alpha key).
"""
torch.manual_seed(0)
w_base = torch.randn(HIDDEN, QKV_OUT)
sd, down, up, magnitude = sd_lora_aitk_magnitude_dora(w_base)
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
mod = next(iter(net.modules.values()))
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (HIDDEN, 1), \
f'dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
updown, _ex_bias = mod.calc_updown(w_base.clone())
merged = w_base + up @ down
norm = merged.reshape(HIDDEN, -1).norm(dim=1).reshape(HIDDEN, 1) + torch.finfo(w_base.dtype).eps
ref = merged * (magnitude.reshape(HIDDEN, 1) / norm) - w_base
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'rel err {rel:.4f}'
return True
def test_lora_magnitude_fused_qkv_sliced():
"""A fused-qkv magnitude vector converts, reshapes and slices per chunk."""
sd = dict(sd_lora_dora_fused_qkv())
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
sd['lora_unet_double_blocks_0_img_attn_qkv.magnitude'] = torch.rand(3 * QKV_OUT) + 0.5
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
for nk, mod in net.modules.items():
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (QKV_OUT, 1), \
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
return True
def test_lora_peft_magnitude_vector():
"""PEFT/diffusers `lora_magnitude_vector` converts the same way."""
sd = {
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
'transformer.transformer_blocks.0.attn.to_q.lora_magnitude_vector': torch.rand(QKV_OUT) + 0.5,
}
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
mod = next(iter(net.modules.values()))
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (QKV_OUT, 1), \
f'dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
return True
def test_lora_sparse_bias_residual():
"""LyCORIS use_bias extraction triplet reconstructs into the dense-bias path.
Mirrors the extraction save exactly: residual sparsified COO with int16
indices, values from the weight-shaped remainder, alpha == rank (scale 1).
Reference: total delta = up @ down + residual.
"""
torch.manual_seed(0)
down = torch.randn(RANK_LORA, QKV_OUT)
up = torch.randn(HIDDEN, RANK_LORA)
residual = torch.randn(HIDDEN, QKV_OUT)
residual[torch.rand_like(residual) < 0.98] = 0.0 # extraction sparsity default
sparse = residual.to_sparse().coalesce()
sd = {
'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': down,
'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': up,
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
'diffusion_model.double_blocks.1.img_attn.proj.bias_indices': sparse.indices().to(torch.int16),
'diffusion_model.double_blocks.1.img_attn.proj.bias_values': sparse.values(),
'diffusion_model.double_blocks.1.img_attn.proj.bias_size': torch.tensor(residual.shape).to(torch.int16),
}
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
mod = next(iter(net.modules.values()))
assert mod.bias is not None and mod.bias.is_sparse, f'bias={type(mod.bias)}'
w_base = torch.randn(HIDDEN, QKV_OUT)
updown, _ex_bias = mod.calc_updown(w_base)
ref = up @ down + residual
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'rel err {rel:.4f}'
return True
def test_lora_sparse_bias_fused_skipped():
"""The sparse residual triplet on a fused target is skipped like dense bias."""
sd = dict(sd_lora_kohya_qkv())
residual = torch.zeros(3 * QKV_OUT, HIDDEN)
residual[0, 0] = 1.0
sparse = residual.to_sparse().coalesce()
base = 'lora_unet_double_blocks_0_img_attn_qkv'
sd[f'{base}.bias_indices'] = sparse.indices().to(torch.int16)
sd[f'{base}.bias_values'] = sparse.values()
sd[f'{base}.bias_size'] = torch.tensor(residual.shape).to(torch.int16)
net = _load_via(F.try_load_lora, sd)
assert net is None, f'expected skip, got {net.modules if net else None}'
return True
def test_lokr_shape_mismatch_rejected():
"""Kron dims that disagree with the module are rejected at load, not at apply."""
sd = sd_lokr_bfl_proj()
# Double the w2 in-dim: kron becomes (QKV_OUT, 2*HIDDEN) vs module (HIDDEN, QKV_OUT).
sd['diffusion_model.double_blocks.1.img_attn.proj.lokr_w2'] = \
torch.randn(QKV_OUT // LOKR_W1_DIM, 2 * (HIDDEN // LOKR_W1_DIM))
net = _load_via(F.try_load_lokr, sd)
assert net is None, f'expected mismatch rejection, got {net.modules if net else None}'
return True
def test_lokr_fused_shape_mismatch_rejected():
"""Fused kron rows must cover total*out; a 2-of-3 sized fused delta is rejected."""
sd = dict(sd_lokr_kohya_qkv())
# Kron rows 2*QKV_OUT instead of 3*QKV_OUT: no valid 3-way equal chunk.
sd['lora_unet_double_blocks_0_img_attn_qkv.lokr_w2'] = \
torch.randn((2 * QKV_OUT) // 8, HIDDEN // 4)
net = _load_via(F.try_load_lokr, sd)
assert net is None, f'expected mismatch rejection, got {net.modules if net else None}'
return True
def test_lokr_bfl_non_fused():
net = _load_via(F.try_load_lokr, sd_lokr_bfl_proj())
assert net is not None and len(net.modules) == 1
@@ -884,6 +1179,28 @@ def test_lokr_simpletuner_lycoris_format():
return True
def test_lokr_bfl_extra_modules_full_matrix():
"""All non-block F2_EXTRA_MAP targets bind; placeholder alpha never scales.
Regression for the gap where resolve_targets only knew single_blocks /
double_blocks and silently dropped embedder / modulation / final-layer
groups as unmapped.
"""
net = _load_via(F.try_load_lokr, sd_lokr_bfl_extra_modules())
bound = len(net.modules) if net else 0
assert net is not None and bound == len(F.F2_EXTRA_MAP), \
f'bound {bound}/{len(F.F2_EXTRA_MAP)}: {sorted(net.modules) if net else []}'
for nk, mod in net.modules.items():
assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk), \
f'{nk}: type={type(mod).__name__}'
# Full-matrix factors leave dim unset; the ~1e10 placeholder alpha must be ignored.
assert mod.dim is None and mod.calc_scale() == 1.0, f'{nk}: dim={mod.dim} scale={mod.calc_scale()}'
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
assert_finite(updown, label=nk)
return True
def test_loha_bfl_non_fused():
net = _load_via(F.try_load_loha, sd_loha_bfl_proj())
assert net is not None and len(net.modules) == 1
@@ -1156,6 +1473,45 @@ def test_full_calc_updown_shape():
return True
def test_dora_ordering_matches_lycoris():
"""DoRA merge equals the LyCORIS forward reference at multiplier 1 and
lerps the full merged delta at other multipliers (ComfyUI semantics).
Reference mirrors lycoris locon forward + apply_weight_decompose with
wd_on_out=True: alpha/rank scales the diff BEFORE the row norms. Uses
alpha != rank (the kohya-style case the old decompose-before-scale
ordering got ~64% wrong).
"""
torch.manual_seed(0)
alpha = 1.0
w_base = torch.randn(QKV_OUT, HIDDEN)
down = torch.randn(RANK_LORA, HIDDEN)
up = torch.randn(QKV_OUT, RANK_LORA)
dora = w_base.reshape(QKV_OUT, -1).norm(dim=1, keepdim=True)
def reference_delta(mult):
diff = (up @ down) * (alpha / RANK_LORA)
merged = w_base + diff
norm = merged.reshape(QKV_OUT, -1).norm(dim=1).reshape(QKV_OUT, 1) + torch.finfo(w_base.dtype).eps
return (merged * (dora / norm) - w_base) * mult
sd = {
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': down,
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': up,
'transformer.transformer_blocks.0.attn.to_q.alpha': torch.tensor(alpha),
'transformer.transformer_blocks.0.attn.to_q.dora_scale': dora,
}
for mult in (1.0, 0.5):
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 1
mod = make_network_for_module(next(iter(net.modules.values())), te_mul=mult, unet_mul=mult)
updown, _ex_bias = mod.calc_updown(w_base.clone())
ref = reference_delta(mult)
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'mult={mult}: rel err {rel:.4f}'
return True
# ============================================================
# Tests — apply path / regressions in shared infra
# ============================================================
@@ -1227,6 +1583,7 @@ def run_tests():
log.warning('=== Parsing primitives ===')
for fn in [test_parse_key_all_prefixes, test_resolve_targets_qkv_chunking,
test_resolve_targets_extra_modules,
test_parse_key_peft_wrapper_unwrap, test_parse_key_lycoris_prefix,
test_parse_key_bare_diffusers_and_peft_default,
test_marker_disambiguation]:
@@ -1243,9 +1600,22 @@ def run_tests():
test_lora_peft_saved_dreambooth_style,
test_lora_peft_saved_diffusers_style,
test_lora_dora_threading,
test_lora_dora_fused_qkv_sliced,
test_lokr_dora_fused_qkv_sliced,
test_dora_per_input_fused_skipped,
test_lora_diff_b_fused_qkv_sliced,
test_legacy_bias_fused_skipped,
test_lora_aitk_magnitude_dora,
test_lora_magnitude_fused_qkv_sliced,
test_lora_peft_magnitude_vector,
test_lora_sparse_bias_residual,
test_lora_sparse_bias_fused_skipped,
test_lokr_shape_mismatch_rejected,
test_lokr_fused_shape_mismatch_rejected,
test_lokr_bfl_non_fused,
test_lokr_kohya_fused_qkv_chunked,
test_lokr_simpletuner_lycoris_format,
test_lokr_bfl_extra_modules_full_matrix,
test_loha_bfl_non_fused,
test_loha_kohya_fused_qkv_chunked,
test_oft_kohya_non_fused,
@@ -1271,6 +1641,7 @@ def run_tests():
test_ia3_calc_updown_shape,
test_glora_calc_updown_shape,
test_full_calc_updown_shape,
test_dora_ordering_matches_lycoris,
]:
run_test(CAT_MATH, fn)
test/test-krea2-transformer.py
+60 -3
View File
@@ -8,8 +8,12 @@
- Zero-init regression: a checkpoint that omits the dormant last.up/last.down residual branch,
after materialize_zero_init, produces the same output as the base whose up is zeroed. Port
only, so it runs without the reference.
- comfy_quant real-file (opt-in): loads an actual ComfyUI int8_tensorwise Krea2 single file
through the native loader and verifies SDNQ adoption plus a finite tiny forward. Enabled by
setting $KREA2_COMFY_FILE to the .safetensors path; needs network for the base repo config
($KREA2_COMFY_REPO, default CalamitousFelicitousness/Krea-2-Base-Diffusers).
No server, no checkpoint.
No server, no checkpoint (except the opt-in comfy_quant test).
"""
import importlib.util
@@ -89,8 +93,8 @@ def run_parity(mmdit, port):
print("PARITY OK")
def load_materialize_zero_init():
"""Import the real loader helper. native_transformer pulls in modules.shared, which needs
def bootstrap_repo():
"""Make repo modules importable. native_transformer pulls in modules.shared, which needs
cmd_args parsed first, so bootstrap it the same way the native-transformer suite does."""
repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if repo not in sys.path:
@@ -106,6 +110,10 @@ def load_materialize_zero_init():
sys.argv = orig_argv
installer.add_args(modules.cmd_args.parser)
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
def load_materialize_zero_init():
bootstrap_repo()
from pipelines.native_transformer import materialize_zero_init
return materialize_zero_init
@@ -159,12 +167,61 @@ def run_zero_init_regression(port):
print("ZERO-INIT OK")
def run_comfy_quant_real_file():
"""Opt-in end-to-end check against a real ComfyUI int8_tensorwise Krea2 file: the native
loader must adopt every marked linear as an SDNQ int8 layer and produce a finite output on
a tiny forward. $KREA2_COMFY_LAYERS overrides the expected layer count (default 224)."""
path = os.environ.get("KREA2_COMFY_FILE")
if not path:
print("COMFY REAL-FILE SKIPPED (set KREA2_COMFY_FILE to enable)")
return
assert os.path.exists(path), f"KREA2_COMFY_FILE not found: {path}"
expected_layers = int(os.environ.get("KREA2_COMFY_LAYERS", "224"))
repo_id = os.environ.get("KREA2_COMFY_REPO", "CalamitousFelicitousness/Krea-2-Base-Diffusers")
bootstrap_repo()
from pipelines import native_transformer as nt
from pipelines.krea2 import KREA2_SPEC
transformer, siblings = nt.load(local_file=path, repo_id=repo_id, spec=KREA2_SPEC, diffusers_cfg={})
assert siblings == {}
sdnq_layers = [m for m in transformer.modules() if m.__class__.__name__ == "SDNQLinear"]
storage_dtypes = {m.weight.dtype for m in sdnq_layers}
print(f"comfy_quant real file: {len(sdnq_layers)} SDNQ layers, storage {storage_dtypes}")
assert len(sdnq_layers) == expected_layers, f"expected {expected_layers} SDNQ layers, got {len(sdnq_layers)}"
assert storage_dtypes <= {torch.int8, torch.float8_e4m3fn, torch.uint8}, f"unexpected storage dtypes: {storage_dtypes}"
assert len(storage_dtypes) == 1, "adopted weights must share one storage dtype"
assert transformer.blocks[0].attn.wq.__class__.__name__ == "SDNQLinear"
assert getattr(transformer, "quantization_config", None) is not None
cfg = transformer.config
param = next(p for p in transformer.parameters() if p.is_floating_point())
device, dtype = param.device, param.dtype
batch, txtlen, imglen = 1, 3, 4
seq = txtlen + imglen
gen = torch.Generator().manual_seed(1)
img = torch.randn(batch, imglen, cfg.channels * cfg.patch ** 2, generator=gen).to(device=device, dtype=dtype)
context = torch.randn(batch, txtlen, cfg.txtlayers, cfg.txtdim, generator=gen).to(device=device, dtype=dtype)
timestep = torch.rand(batch, generator=gen).to(device=device, dtype=dtype)
pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float().to(device=device)
mask = torch.ones(batch, seq, dtype=torch.bool, device=device)
with torch.no_grad():
out = transformer(
hidden_states=img, encoder_hidden_states=context, timestep=timestep,
position_ids=pos, attention_mask=mask, return_dict=False,
)[0]
assert torch.isfinite(out).all(), "forward produced non-finite values"
print("COMFY REAL-FILE OK")
def main():
port = load_port()
# Parity runs first in pristine torch state; the regression imports modules afterwards.
mmdit = load_reference()
run_parity(mmdit, port)
run_zero_init_regression(port)
run_comfy_quant_real_file()
if __name__ == "__main__":
test/test-native-transformer.py
+1530 -2
View File
File diff suppressed because it is too large Load Diff
test/test-zimage-native-adapters.py
+151 -1
View File
@@ -125,6 +125,7 @@ MLP_HIDDEN = int(HIDDEN / 3 * 8) # 256, matches ZImageTransformerBlock FeedFor
ADALN_OUT = 4 * HIDDEN # 384, matches Z-Image adaLN_modulation
N_LAYERS = 2 # main transformer blocks
N_REFINER = 1 # noise_refiner / context_refiner blocks
PATCH_KEY = "2-1" # the single all_x_embedder / all_final_layer key both variants build
# pylint: disable=attribute-defined-outside-init
@@ -184,6 +185,26 @@ def build_mock_transformer():
transformer.layers = torch.nn.ModuleList([build_zimage_block(modulation=False) for _ in range(N_LAYERS)])
transformer.noise_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)])
transformer.context_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)])
# Non-block targets. all_x_embedder / all_final_layer are ModuleDicts keyed by
# "{patch_size}-{f_patch_size}"; both shipped variants build the single key below.
final_layer = _Holder()
final_layer.linear = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
final_layer.adaLN_modulation = torch.nn.Sequential(
torch.nn.SiLU(),
torch.nn.Linear(HIDDEN, ADALN_OUT, bias=True),
)
transformer.all_x_embedder = torch.nn.ModuleDict({PATCH_KEY: torch.nn.Linear(HIDDEN, HIDDEN, bias=True)})
transformer.all_final_layer = torch.nn.ModuleDict({PATCH_KEY: final_layer})
transformer.t_embedder = _Holder()
transformer.t_embedder.mlp = torch.nn.Sequential(
torch.nn.Linear(HIDDEN, HIDDEN, bias=True),
torch.nn.SiLU(),
torch.nn.Linear(HIDDEN, HIDDEN, bias=True),
)
transformer.cap_embedder = torch.nn.Sequential(
torch.nn.RMSNorm(HIDDEN),
torch.nn.Linear(HIDDEN, HIDDEN, bias=True),
)
return transformer
@@ -313,6 +334,24 @@ def sd_lokr_legacy_fused_qkv():
}
def sd_lokr_lycoris_style():
"""LyCORIS-standalone LoKR (``lycoris_`` + underscored diffusers path).
The prefix is resolved by the universal passthrough in native_adapter, not
by anything zimage-specific. ``to_out.0`` exercises the verbatim
round-trip where a naive underscore-to-dot expansion would corrupt the
ModuleList index.
"""
return {
'lycoris_layers_0_attention_to_q.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
'lycoris_layers_0_attention_to_q.lokr_w2': torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
'lycoris_layers_0_attention_to_q.alpha': torch.tensor(float(LOKR_W1_DIM)),
'lycoris_layers_1_attention_to_out_0.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
'lycoris_layers_1_attention_to_out_0.lokr_w2': torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
'lycoris_layers_1_attention_to_out_0.alpha': torch.tensor(float(LOKR_W1_DIM)),
}
def sd_loha_bfl_proj():
"""LoHA on attention.to_out.0 (non-fused; LoHA on fused qkv is skipped by the loader)."""
return {
@@ -599,6 +638,112 @@ def test_lokr_legacy_fused_qkv_chunked():
return True
def test_lokr_lycoris_prefix_passthrough():
"""lycoris_ keys load via the universal passthrough (hoisted, not zimage-specific)."""
net = _load_via(Z.try_load_lokr, sd_lokr_lycoris_style())
assert net is not None and len(net.modules) == 2, f'got {net.modules if net else None}'
expected = {
'lora_transformer_layers_0_attention_to_q',
'lora_transformer_layers_1_attention_to_out_0',
}
assert set(net.modules) == expected, f'got {set(net.modules)}'
for mod in net.modules.values():
assert isinstance(mod, network_lokr.NetworkModuleLokr)
return True
def test_full_diff_chain():
"""Full-diff extraction loads through the chain (RedLSP-shaped).
Verifies the try_load chain includes try_load_full: diff+diff_b on the
legacy attention.out alias binds; the fused attention.qkv diff group is
skipped with a warning (no chunk variant for the Full family).
"""
sd = {
'diffusion_model.layers.0.attention.out.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.layers.0.attention.out.diff_b': torch.randn(HIDDEN),
'diffusion_model.layers.0.attention.qkv.diff': torch.randn(3 * HIDDEN, HIDDEN),
}
net = _load_via(Z.try_load, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
assert 'lora_transformer_layers_0_attention_to_out_0' in net.modules, f'got {set(net.modules)}'
mod = next(iter(net.modules.values()))
updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
assert tuple(updown.shape) == (HIDDEN, HIDDEN) and torch.isfinite(updown).all()
assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,)
return True
def test_full_stamps_norm_targets():
"""A Full adapter on a norm target binds AND applies.
assign_network_names_to_compvis_modules puts norms in the mapping but never
stamps network_layer_name on them, and the apply pass keys off that
attribute: without loader-local stamping the module binds and then silently
never applies. Regression for the RedLSP extraction, which bound 308
modules and applied only 172.
"""
sd = {
'diffusion_model.layers.0.attention_norm1.diff': torch.randn(HIDDEN),
'diffusion_model.layers.0.attention.q_norm.diff': torch.randn(HEAD_DIM),
}
net = _load_via(Z.try_load, sd)
assert net is not None and len(net.modules) == 2, f'got {sorted(net.modules) if net else None}'
expected = {'lora_transformer_layers_0_attention_norm1', 'lora_transformer_layers_0_attention_norm_q'}
assert set(net.modules) == expected, f'got {set(net.modules)}'
for key, mod in net.modules.items():
assert getattr(mod.sd_module, 'network_layer_name', None) == key, \
f'{key}: host module not stamped, the apply pass will skip it'
return True
def test_resolve_targets_norm_aliases_and_extra():
"""qk-norm renames and the ModuleDict-keyed non-block targets."""
cases = [
(('diffusion_model.', 'layers.0.attention.q_norm'), 'layers.0.attention.norm_q'),
(('diffusion_model.', 'layers.3.attention.k_norm'), 'layers.3.attention.norm_k'),
(('lora_unet_', 'layers_0_attention_q_norm'), 'layers_0_attention_norm_q'),
((None, 'noise_refiner.0.attention.k_norm'), 'noise_refiner.0.attention.norm_k'),
]
for (prefix, base), expected in cases:
targets = Z.resolve_targets(prefix, base)
assert targets == [(expected, None)], f'({prefix}, {base}) -> {targets}'
# the patch key comes from the live model's all_x_embedder ModuleDict
install_mock_pipe()
for base, expected in [
('x_embedder', f'all_x_embedder.{PATCH_KEY}'),
('final_layer.linear', f'all_final_layer.{PATCH_KEY}.linear'),
('final_layer.adaLN_modulation.1', f'all_final_layer.{PATCH_KEY}.adaLN_modulation.1'),
]:
targets = Z.resolve_targets('diffusion_model.', base)
assert targets == [(expected, None)], f'{base} -> {targets}'
kohya = Z.resolve_targets('lora_unet_', base.replace('.', '_'))
assert kohya == [(expected, None)], f'kohya {base} -> {kohya}'
return True
def test_full_extra_modules_bind():
"""x_embedder and final_layer bind through the ModuleDict paths."""
sd = {
'diffusion_model.x_embedder.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.final_layer.linear.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.final_layer.adaLN_modulation.1.diff': torch.randn(ADALN_OUT, HIDDEN),
'diffusion_model.t_embedder.mlp.0.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.cap_embedder.1.diff': torch.randn(HIDDEN, HIDDEN),
}
net = _load_via(Z.try_load, sd)
assert net is not None and len(net.modules) == 5, f'got {sorted(net.modules) if net else None}'
expected = {
f'lora_transformer_all_x_embedder_{PATCH_KEY}',
f'lora_transformer_all_final_layer_{PATCH_KEY}_linear',
f'lora_transformer_all_final_layer_{PATCH_KEY}_adaLN_modulation_1',
'lora_transformer_t_embedder_mlp_0',
'lora_transformer_cap_embedder_1',
}
assert set(net.modules) == expected, f'got {set(net.modules)}'
return True
def test_loha_bfl_proj():
"""BFL LoHA on a non-fused proj target binds via NetworkModuleHada."""
net = _load_via(Z.try_load_loha, sd_loha_bfl_proj())
@@ -714,7 +859,8 @@ def run_tests():
t0 = time.time()
log.warning('=== Parsing primitives ===')
for fn in [test_parse_key_all_prefixes, test_marker_disambiguation]:
for fn in [test_parse_key_all_prefixes, test_marker_disambiguation,
test_resolve_targets_norm_aliases_and_extra]:
run_test(CAT_PARSE, fn)
log.warning('=== Loaders ===')
@@ -729,6 +875,10 @@ def run_tests():
test_lora_dora_threading,
test_lokr_bfl_adaln,
test_lokr_legacy_fused_qkv_chunked,
test_lokr_lycoris_prefix_passthrough,
test_full_diff_chain,
test_full_stamps_norm_targets,
test_full_extra_modules_bind,
test_loha_bfl_proj,
test_loha_legacy_fused_qkv_chunked,
test_oft_lycoris_no_npe,
ui/authWrap.ts
+6
View File
@@ -7,6 +7,7 @@ interface TokenResponse {
let user: string | undefined;
let token: string | undefined;
let baseURL: string | undefined;
export async function getToken(): Promise<{ user: string | undefined; token: string | undefined }> {
if (token === undefined || user === undefined) {
@@ -38,6 +39,11 @@ export async function authFetch(url: RequestInfo | URL, options: RequestInit = {
error('fetch', { status: res?.status || 503, url, user, token, error: err });
}
}
if (!baseURL) {
// baseURL = `${window.location.protocol}//${window.location.host}:${window.location.port}`;
baseURL = window.location.origin;
log('origin', baseURL);
}
return res;
}
window.authFetch = authFetch;
ui/css/sdnext.css
+11 -1
View File
@@ -1852,7 +1852,7 @@ background: var(--background-color)
#tab-gallery-files gallery-file {
/* Add a vertical gutter between items (left/right), matching existing small row spacing */
display: inline-block !important;
display: inline-block;
margin-right: 0.2em;
vertical-align: top; /* keep rows aligned on the top edge */
}
@@ -2286,6 +2286,16 @@ div:has(>#tab-gallery-folders) {
filter: blur(0);
}
#logMonitor {
width: stretch;
}
#logMonitorData {
height: 12vh;
overflow-y: scroll;
overflow-x: auto;
}
.disabled-look textarea {
opacity: 0.4;
pointer-events: none;
ui/dist/sdnext.mjs
Vendored
+144 -84
View File
@@ -9883,6 +9883,7 @@ window.xhrPost = xhrPost;
// ui/authWrap.ts
var user;
var token;
var baseURL;
async function getToken() {
if (token === void 0 || user === void 0) {
const res = await fetch(`${window.subpath}/token`);
@@ -9912,6 +9913,10 @@ async function authFetch(url2, options = {}) {
error("fetch", { status: res?.status || 503, url: url2, user, token, error: err });
}
}
if (!baseURL) {
baseURL = window.location.origin;
log("origin", baseURL);
}
return res;
}
window.authFetch = authFetch;
@@ -10152,7 +10157,7 @@ async function initTableSorter() {
async function deleteFile(filename) {
if (!filename) return;
if (!confirm(`Are you sure you want to delete the object - This action cannot be undone? Object: ${filename}`)) return;
const res = await authFetch(`${window.api}/delete-file?file=${encodeURIComponent(filename)}`);
const res = await authFetch(`${window.api}/delete-file?file=${encodeURIComponent(filename)}`, { method: "DELETE" });
if (!res || res.status !== 200) {
error("FileDelete", { file: filename, status: res?.status, statusText: res?.statusText });
return;
@@ -10308,9 +10313,13 @@ async function setupControlUI() {
// ui/extraNetworks.ts
var activePromptTextarea = {};
var selectedNetworks = {};
var sortVal = -1;
var totalCards = -1;
var lastTab = "control";
function getSelectedNetworks() {
return selectedNetworks;
}
var getENActiveTab = () => {
let tabName = "";
if (gradioApp().getElementById("txt2img_prompt")?.checkVisibility() || gradioApp().getElementById("txt2img_generate")?.checkVisibility()) tabName = "txt2img";
@@ -10539,6 +10548,7 @@ function sortExtraNetworks(fixed = "no") {
}
async function markSelectedCards(selected, page = "") {
log("markSelectedCards", selected, page);
selectedNetworks[page] = selected;
gradioApp().querySelectorAll(".extra-network-cards .card").forEach((el2) => {
if (page.length > 0 && el2.dataset.page !== page) return;
if (selected.includes(el2.dataset.name) || selected.includes(el2.dataset.short)) el2.classList.add("card-selected");
@@ -10851,6 +10861,7 @@ window.applyStyles = applyStyles;
window.closeDetailsEN = closeDetailsEN;
window.getENActivePage = getENActivePage;
window.getCardDetails = getCardDetails;
window.getSelectedNetworks = getSelectedNetworks;
window.sortExtraNetworks = sortExtraNetworks;
window.refeshDetailsEN = refeshDetailsEN;
window.extraNetworksSearchButton = extraNetworksSearchButton;
@@ -10965,7 +10976,7 @@ var progressTimeout = 180;
var startTimeout = 5;
function setRefreshInterval() {
refreshInterval = window.opts.live_preview_refresh_period || 500;
log("refreshInterval", document.visibilityState, refreshInterval);
log("refreshInterval", { visibile: document.visibilityState, interval: refreshInterval });
document.addEventListener("visibilitychange", () => {
if (window.opts.live_preview_require_focus !== false && document.hidden) refreshInterval = Math.max(2500, window.opts.live_preview_refresh_period || 1e3);
else refreshInterval = window.opts.live_preview_refresh_period || 1e3;
@@ -11287,14 +11298,13 @@ function setFontSize(val, old) {
rootStyle.setProperty("--text-xxl", `${nextSize + 3}px`);
appliedFontSize = nextSize;
const t1 = performance.now();
log("setFontSize", nextSize, `time=${Math.round(t1 - t0)}`);
timer("setFontSize", t1 - t0);
log("setFontSize", { size: nextSize, time: Math.round(t1 - t0) });
});
}
function switchToTab(tab) {
const tabs = Array.from(gradioApp().querySelectorAll("#tabs > .tab-nav > button"));
const btn = tabs?.find((t) => t.innerText === tab);
log("switchToTab", tab);
log("switchToTab", { tab });
if (btn) btn.click();
}
function switch_to_txt2img(...args) {
@@ -11662,8 +11672,12 @@ function getDesiredCheckpointName() {
}
function selectUNet(name) {
desiredUNetName = name;
gradioApp().getElementById("change_unet").click();
log(`selectUNet: ${desiredUNetName}`);
const tabName = getENActiveTab();
const btnModel = gradioApp().getElementById(`${tabName}_extra_model`);
const isSecondary = btnModel && btnModel.classList.contains("toolbutton-selected");
if (isSecondary) gradioApp().getElementById("change_unet_secondary").click();
else gradioApp().getElementById("change_unet").click();
log(`selectUNet ${isSecondary ? "secondary" : "primary"}: ${desiredUNetName}`);
markSelectedCards([desiredUNetName], "unet");
}
function selectReference(name) {
@@ -11988,14 +12002,11 @@ var logErrors = 0;
var logConnected = false;
function dateToStr(ts) {
const dt = new Date(1e3 * ts);
const year = dt.getFullYear();
const mo = String(dt.getMonth() + 1).padStart(2, "0");
const day = String(dt.getDate()).padStart(2, "0");
const hour = String(dt.getHours()).padStart(2, "0");
const min = String(dt.getMinutes()).padStart(2, "0");
const sec = String(dt.getSeconds()).padStart(2, "0");
const ms = String(dt.getMilliseconds()).padStart(3, "0");
const s = `${year}-${mo}-${day} ${hour}:${min}:${sec}.${ms}`;
const s = `${hour}:${min}:${sec}.${ms}`;
return s;
}
function htmlEscape(text) {
@@ -12014,6 +12025,21 @@ function parseLogLine(line) {
msg: String(parsed.msg ?? "")
};
}
async function clearErrors() {
logWarnings = 0;
logErrors = 0;
log("clearErrors");
}
async function initClearErrorsButton() {
const btnServerClear = document.getElementById("btn_console_log_server_clear");
if (btnServerClear) {
btnServerClear.onclick = async (evt) => {
evt.preventDefault();
evt.stopPropagation();
clearErrors();
};
}
}
async function logMonitor() {
const addLogLine = (line) => {
if (!logMonitorEl) logMonitorEl = document.getElementById("logMonitorData");
@@ -12024,14 +12050,13 @@ async function logMonitor() {
const level = `<td style="color: var(--color-${l.level.toLowerCase()})">${l.level}</td>`;
if (l.level === "WARNING") logWarnings++;
if (l.level === "ERROR") logErrors++;
const module = `<td style="color: var(--neutral-400)">${l.module}</td>`;
const module = `<td style="color: #ffca68">${l.module}</td>`;
const facilityText = l.facility.length > 20 ? `${l.facility.substring(0, 20)}...` : l.facility;
const facility = l.facility !== "sd" ? `<td>${facilityText}</td>` : "<td></td>";
const facility = l.facility !== "sd" ? `<td style="color: #ffca68">${facilityText}</td>` : "<td></td>";
row.innerHTML = `<td>${dateToStr(l.created)}</td>${level}${facility}${module}<td>${htmlEscape(l.msg)}</td>`;
logMonitorEl.appendChild(row);
} catch (err) {
error(`logMonitor: ${String(err)}`);
error(`logMonitor: ${line}`);
error("logMonitor", { error: String(err), line });
}
};
const cleanupLog = (atBottom2) => {
@@ -12046,7 +12071,13 @@ async function logMonitor() {
const modenUIBtn = document.getElementById("btn_console");
if (elWarn) elWarn.innerText = String(logWarnings);
if (elErr) elErr.innerText = String(logErrors);
if (modenUIBtn) modenUIBtn.setAttribute("error-count", logErrors > 0 ? String(logErrors) : "");
if (modenUIBtn) {
modenUIBtn.setAttribute("error-count", logErrors > 0 ? String(logErrors) : "");
modenUIBtn.style.backgroundColor = logErrors > 0 ? "var(--color-error)" : "";
modenUIBtn.title = `Log
Errors ${logErrors}
Warnings ${logWarnings}`;
}
};
const txtGallery = document.getElementById("txt2img_gallery");
if (txtGallery) txtGallery.style.height = window.opts.logmonitor_show ? "50vh" : "55vh";
@@ -12080,51 +12111,66 @@ async function logMonitor() {
if (logMonitorEl && lines?.length > 0 && logMonitorEl.parentElement?.parentElement instanceof HTMLElement) {
logMonitorEl.parentElement.parentElement.style.display = window.opts.logmonitor_show ? "block" : "none";
}
for (const line of lines) addLogLine(line);
if (!logConnected) {
logConnected = true;
xhrPost(`${window.api}/log`, { debug: "connected" });
logErrors = 0;
}
for (const line of lines) addLogLine(line);
} else {
logConnected = false;
logErrors++;
addLogLine(`{ "created": ${Date.now()}, "level":"ERROR", "module":"logMonitor", "facility":"ui", "msg":"Failed to fetch log: ${res?.status} ${res?.statusText}" }`);
if (res) addLogLine(`{ "created": ${Date.now()}, "level":"ERROR", "module":"logMonitor", "facility":"ui", "msg":"Failed to fetch log: ${res?.status} ${res?.statusText}" }`);
else addLogLine(`{ "created": ${Date.now()}, "level":"ERROR", "module":"logMonitor", "facility":"ui", "msg":"Server unreachable" }`);
}
cleanupLog(atBottom);
} catch {
logConnected = false;
logErrors++;
addLogLine(`{ "created": ${Date.now()}, "level":"ERROR", "module":"logMonitor", "facility":"ui", "msg":"Failed to fetch log: server unreachable" }`);
addLogLine(`{ "created": ${Date.now()}, "level":"ERROR", "module":"logMonitor", "facility":"ui", "msg":"Server unreachable" }`);
cleanupLog(atBottom);
}
}
async function initLogMonitor() {
const el2 = document.getElementsByTagName("footer")[0];
let el2 = document.getElementById("logMonitorPlaceholder");
const modernUi = Boolean(el2);
if (!el2) el2 = document.getElementsByTagName("footer")[0];
if (!el2) return;
const t0 = performance.now();
el2.classList.add("log-monitor");
const uiDisabled = Array.isArray(window.opts.ui_disabled) ? window.opts.ui_disabled : [];
if (uiDisabled.includes("logs")) return;
el2.innerHTML = `
<table id="logMonitor" style="width: 100%;">
<thead style="display: block; text-align: left; border-bottom: solid 1px var(--button-primary-border-color)">
<tr>
<th style="width: 144px">Time</th>
<th>Level</th>
<th style="width: 0"></th>
<th style="width: 154px">Module</th>
<th>Message</th>
<th style="position: absolute; right: 7em">Warnings <span id="logWarnings">0</span></th>
<th style="position: absolute; right: 1em">Errors <span id="logErrors">0</span></th>
</tr>
</thead>
<tbody id="logMonitorData" style="white-space: nowrap; height: 10vh; width: 100vw; display: block; overflow-x: hidden; overflow-y: scroll; color: var(--neutral-400)">
</tbody>
</table>
`;
if (modernUi) {
el2.style.overflow = "auto";
el2.innerHTML = `
<table id="logMonitor" style="width: 100%;">
<tbody id="logMonitorData" style="white-space: nowrap; display: block">
</tbody>
</table>
`;
} else {
el2.innerHTML = `
<table id="logMonitor" style="width: 100%;">
<thead style="display: block; text-align: left; border-bottom: solid 1px var(--button-primary-border-color)">
<tr>
<th style="width: 144px">Time</th>
<th>Level</th>
<th style="width: 0"></th>
<th style="width: 154px">Module</th>
<th>Message</th>
<th style="position: absolute; right: 7em">Warnings <span id="logWarnings">0</span></th>
<th style="position: absolute; right: 1em">Errors <span id="logErrors">0</span></th>
</tr>
</thead>
<tbody id="logMonitorData" style="white-space: nowrap; display: block; color: var(--neutral-400)">
</tbody>
</table>
`;
}
el2.style.display = "none";
authFetch(`${window.api}/start?agent=${encodeURI(navigator.userAgent)}`);
logMonitor();
initClearErrorsButton();
const t1 = performance.now();
log("initLogMonitor", { show: window.opts.logmonitor_show, time: Math.round(t1 - t0) });
timer("initLogMonitor", t1 - t0);
@@ -12183,7 +12229,7 @@ async function updateOpts(json_string) {
else opts_tabs[meta.tab_name].saved_keys.add(opt);
});
const t2 = performance.now();
log("updateOpts", `settings=${Object.keys(new_opts).length} callbacks=${Math.round(t2 - t1)} apply=${Math.round(t1 - t0)}`);
log("updateOpts", { settings: Object.keys(new_opts).length, callbacks: Math.round(t2 - t1), apply: Math.round(t1 - t0) });
timer("updateOpts", t2 - t0);
}
function showAllSettings() {
@@ -12926,7 +12972,6 @@ var minCleanupCount = 1e3;
var minCleanupTime = 1e3 * 60 * 60;
var folderStylesheet = new CSSStyleSheet();
var fileStylesheet = new CSSStyleSheet();
var iconStopwatch = String.fromCodePoint(9201);
var separatorStates = /* @__PURE__ */ new Map();
var el = {
folders: void 0,
@@ -12940,6 +12985,13 @@ var el = {
var cleanupTimers = {};
var maintenanceTimers = {};
var fetchQueue = [];
var icons = {
Time: String.fromCodePoint(9201),
Folder: String.fromCodePoint(128448),
// or 128449;
Sort: String.fromCodePoint(8645),
Images: String.fromCodePoint(128461)
};
var SUPPORTED_EXTENSIONS = ["jpg", "jpeg", "png", "webp", "tiff", "jp2", "jxl", "gif", "mp4", "mkv", "avi", "mjpeg", "mpg", "avr"];
var gallerySorter = {
nameA: { name: "Name Ascending", func: (a, b) => a.name.localeCompare(b.name) },
@@ -13099,9 +13151,9 @@ function updateGalleryStyles() {
padding: 4px;
font-size: 1.2em;
letter-spacing: 0.5em;
width: 140px;
margin-top: calc(140px - 32px);
opacity: 75%;
border-radius: var(--sd-border-radius);
}
:host(.gallery-file-selected) .gallery-file {
box-shadow: 0 0 0 2px var(--sd-button-selected-color);
@@ -13151,6 +13203,10 @@ var SimpleProgressBar = class {
if (total <= 0) return;
this.hide();
this.#max = total;
if (this.#monitoredSet.size >= this.#max) {
this.stop();
return;
}
this.#interval = setInterval(() => this.update(this.#monitoredSet.size, this.#max), 100);
}
attachTo(element) {
@@ -13174,7 +13230,6 @@ var SimpleProgressBar = class {
}
stop() {
clearInterval(this.#interval);
this.#interval = void 0;
if (this.stats.count) {
debug("gallery: thumbnail stats", this.stats);
this.stats = { ...this.defaultStats };
@@ -13461,7 +13516,7 @@ Size: ${this.size.toLocaleString()} bytes
Modified: ${this.mtime.toLocaleString()}`;
this.title = img.title;
const shouldDisplayBasedOnSearch = this.title.toLowerCase().includes(el.search.value.toLowerCase());
if (this.style.display !== "none") this.style.display = shouldDisplayBasedOnSearch ? "flex" : "none";
if (this.style.display !== "none") this.style.display = shouldDisplayBasedOnSearch ? "" : "none";
this.shadow.appendChild(img);
pb.stats.elapsed = (pb.stats.elapsed || 0) + Math.round(performance.now() - t0);
}
@@ -13478,7 +13533,7 @@ async function handleSeparator(separator) {
const fileDir = f.name.match(/(.*)[/\\]/);
const fileDirPath = fileDir ? fileDir[1] : "";
if (separator.title.length > 0 && fileDirPath === separator.title) {
f.style.display = nowHidden ? "none" : "unset";
f.style.display = nowHidden ? "none" : "";
}
}
}
@@ -13547,24 +13602,22 @@ async function addSeparators() {
}
var gallerySendImage = (_images) => [currentImage];
window.gallerySendImage = gallerySendImage;
function updateStatusWithSort(...messages) {
function updateStatusLine(...messages) {
if (!el.status) return;
messages.unshift(["Sort", sortMode.name]);
const fragment = document.createDocumentFragment();
for (let i = 0; i < messages.length; i++) {
const div = document.createElement("div");
if (Array.isArray(messages[i])) {
const [text1, text2] = messages[i];
const tDiv1 = document.createElement("div");
tDiv1.innerText = `${text1}:`;
const tDiv2 = document.createElement("div");
tDiv2.innerText = text2;
tDiv2.title = text2;
div.append(tDiv1, tDiv2);
const [k, v] = messages[i];
const tDiv = document.createElement("div");
const ico = icons[k] || `${k}:`;
tDiv.innerText = `${ico} ${v}`;
div.append(tDiv);
} else {
const tDiv1 = document.createElement("div");
tDiv1.innerText = messages[i];
div.append(tDiv1);
const tDiv = document.createElement("div");
tDiv.innerText = messages[i];
div.append(tDiv);
}
fragment.append(div);
}
@@ -13643,7 +13696,7 @@ async function gallerySearch() {
const isOpen = separatorStates.get(dirPath);
f.style.display = !dirPath || isOpen ? "unset" : "none";
});
updateStatusWithSort("Filter", "Cleared", ["Images", allFiles.length.toLocaleString()]);
updateStatusLine("Filter", "Cleared", ["Images", allFiles.length.toLocaleString()]);
return;
}
let totalFound = 0;
@@ -13688,7 +13741,7 @@ async function gallerySearch() {
f.style.display = fileMatches.has(f) ? "unset" : "none";
}
const t1 = performance.now();
updateStatusWithSort("Filter", ["Images", `${totalFound.toLocaleString()} / ${allFiles.length.toLocaleString()}`], `${iconStopwatch} ${Math.round(t1 - t0).toLocaleString()}ms`);
updateStatusLine("Filter", ["Images", `${totalFound.toLocaleString()} / ${allFiles.length.toLocaleString()}`], ["Time", `${Math.round(t1 - t0).toLocaleString()}ms`]);
timer(`galleryFilter:${str}`, t1 - t0);
refreshGallerySelection();
}, 250);
@@ -13713,15 +13766,15 @@ async function gallerySort(key) {
const folderGroups = /* @__PURE__ */ new Map();
for (const file of subfolderFiles) {
const dir = getDirPath(file);
if (!folderGroups.has(dir)) {
folderGroups.set(dir, []);
}
if (!folderGroups.has(dir)) folderGroups.set(dir, []);
folderGroups.get(dir).push(file);
}
sortMode = gallerySorter[currentSort];
rootFiles.sort(sortMode.func);
rootFiles.forEach((node) => fragment.appendChild(node));
const sortedFolderNames = Array.from(folderGroups.keys()).sort((a, b) => a.localeCompare(b));
const folderNames = Array.from(folderGroups.keys());
const sortedFolderNames = currentSort.endsWith("A") ? folderNames.sort((a, b) => a.localeCompare(b)) : folderNames.sort((a, b) => b.localeCompare(a));
console.log("HERE", sortedFolderNames);
for (const folderName of sortedFolderNames) {
const files = folderGroups.get(folderName);
files.sort(sortMode.func);
@@ -13745,7 +13798,7 @@ async function gallerySort(key) {
}
const t1 = performance.now();
log(`gallerySort: sort=${sortMode.name} len=${arr.length} time=${Math.floor(t1 - t0)}`);
updateStatusWithSort(["Images", arr.length.toLocaleString()], `${iconStopwatch} ${Math.round(t1 - t0).toLocaleString()}ms`);
updateStatusLine(["Images", arr.length.toLocaleString()], ["Time", `${Math.round(t1 - t0).toLocaleString()}ms`]);
timer(`gallerySort:${sortMode.name}`, t1 - t0);
refreshGallerySelection();
}
@@ -13815,7 +13868,7 @@ async function thumbCacheCleanup(folder, imgCount, controller, force = false) {
log("galleryMaintenance", { folder, kept: keptGalleryHashes.size, deleted: delcount, time: Math.round(t1 - t0) });
timer(`thumbnailDBCleanup:${folder}`, t1 - t0);
currentGalleryFolder = null;
updateStatusWithSort("Thumbnail cache cleared");
updateStatusLine("Thumbnail cache cleared");
}).catch((reason) => {
SimpleFunctionQueue.abortLogger("thumbCacheCleanup", reason);
}).finally(async () => {
@@ -13847,11 +13900,11 @@ window.clearCache = clearCache;
async function fetchFilesHT(evt, controller) {
const t0 = performance.now();
const fragment = document.createDocumentFragment();
updateStatusWithSort(["Folder", evt.target.name], "in-progress");
updateStatusLine(["Folder", evt.target.name], "in-progress");
let numFiles = 0;
const res = await authFetch(`${window.api}/browser/files?folder=${encodeURI(evt.target.name)}`);
if (!res || res.status !== 200) {
updateStatusWithSort(["Folder", evt.target.name], ["Failed", res?.statusText || "No response"]);
updateStatusLine(["Folder", evt.target.name], ["Failed", res?.statusText || "No response"]);
return;
}
const jsonData = await res.json();
@@ -13870,7 +13923,7 @@ async function fetchFilesHT(evt, controller) {
const t1 = performance.now();
log(`gallery: folder=${evt.target.name} num=${numFiles} method=http time=${Math.floor(t1 - t0)}ms`);
timer(`galleryFetch:${evt.target.name}`, t1 - t0);
updateStatusWithSort(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], ["Time", `${Math.floor(t1 - t0).toLocaleString()}ms`]);
pb.start(numFiles);
addSeparators();
refreshGallerySelection();
@@ -13896,7 +13949,7 @@ async function fetchFilesWS(evt) {
await fetchFilesHT(evt, controller);
return;
}
updateStatusWithSort(["Folder", evt.target.name]);
updateStatusLine(["Folder", evt.target.name]);
const t0 = performance.now();
let numFiles = 0;
let t1 = performance.now();
@@ -13914,7 +13967,7 @@ async function fetchFilesWS(evt) {
numFiles++;
fragment.appendChild(file);
if (numFiles % fragmentSize === 0) {
updateStatusWithSort(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], "in-progress", `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], ["Status", "in-progress"], ["Time", `${Math.floor(t1 - t0).toLocaleString()}ms`]);
el.files.appendChild(fragment);
fragment = document.createDocumentFragment();
}
@@ -13925,7 +13978,7 @@ async function fetchFilesWS(evt) {
if (controller.signal.aborted) return;
el.files.appendChild(fragment);
log(`gallery: folder=${evt.target.name} num=${numFiles} method=ws time=${Math.floor(t1 - t0)}ms`);
updateStatusWithSort(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(["Folder", evt.target.name], ["Images", numFiles.toLocaleString()], ["Time", `${Math.floor(t1 - t0).toLocaleString()}ms`]);
pb.start(numFiles);
addSeparators();
refreshGallerySelection();
@@ -14011,7 +14064,7 @@ async function initGalleryAutoRefresh() {
galleryVisObserver.observe(galleryTab, { attributeFilter: ["class", "style"], attributeOldValue: true });
}
async function overlayDelete(evt) {
const res = await authFetch(`${window.api}/delete-image?file=${encodeURIComponent(currentImage)}`);
const res = await authFetch(`${window.api}/delete-image?file=${encodeURIComponent(currentImage)}`, { method: "DELETE" });
evt.stopPropagation();
if (!res || res.status !== 200) {
error("galleryDelete", { file: currentImage, status: res?.status, statusText: res?.statusText });
@@ -15427,7 +15480,7 @@ async function getLocaleData(desiredLocale = null) {
localeData.locale = desiredLocale || "en";
localeData.prev = localeData.locale;
}
log("getLocale", desiredLocale, localeData.locale);
log("getLocale", { lang: desiredLocale, locale: localeData.locale });
let json = {};
try {
let res = await fetch(`${window.subpath}/file=ui/locale/locale_${localeData.locale}.json`);
@@ -15873,7 +15926,10 @@ async function createSplash() {
monitorLog();
await authFetch(`${window.api}/motd`).then((res) => res.text()).then((text) => {
const clean = text.replace(/["]+/g, "");
log("getMOTD", clean);
const boldMatch = clean.match(/<b>(.*?)<\/b>/);
const boldText = boldMatch ? boldMatch[1] : clean;
if (boldMatch) log("getMOTD", { version: boldText });
else log("getMOTD", { text: clean });
const motdEl = document.getElementById("motd");
if (motdEl) motdEl.innerHTML = clean;
}).catch((err) => error(`getMOTD: ${err}`));
@@ -15922,6 +15978,14 @@ async function postStartup() {
disconnectHintsObserver();
logTimers();
}
async function updateSubpath() {
log("mountURL", window.opts.subpath);
if (window.opts.subpath?.length > 0) {
window.subpath = window.opts.subpath;
window.api = `${window.subpath}/sdapi/v1`;
}
log("API", { url: window.api });
}
async function initStartup() {
const t0 = performance.now();
log("initGradio", Math.round(t0 - appStartTime));
@@ -15942,14 +16006,10 @@ async function initStartup() {
startupPromises.push(Promise.resolve(setupControlUI()));
await reconnectUI();
await waitForOpts();
log("mountURL", window.opts.subpath);
if (window.opts.subpath?.length > 0) {
window.subpath = window.opts.subpath;
window.api = `${window.subpath}/sdapi/v1`;
}
startupPromises.push(initLogMonitor());
await updateSubpath();
executeCallbacks(uiReadyCallbacks);
if (window.waitForUiReady) await window.waitForUiReady();
startupPromises.push(Promise.resolve(initLogMonitor()));
startupPromises.push(Promise.resolve(initGallery()));
startupPromises.push(Promise.resolve(setRefreshInterval()));
startupPromises.push(Promise.resolve(setupExtraNetworks()));
@@ -19127,12 +19187,12 @@ async function updateGPU() {
}
const gpuTbody = gpuTable.querySelector("tbody");
if (!gpuTbody) return;
for (const gpu of data) {
let rows = `<tr><td>GPU</td><td>${gpu.name}</td></tr>`;
for (const item of Object.entries(gpu.data)) rows += `<tr><td>${item[0]}</td><td>${item[1]}</td></tr>`;
gpuTbody.innerHTML = rows;
if (gpu.chart && gpu.chart.length === 2) updateGPUChart(gpu.chart[0], gpu.chart[1]);
}
let gpu = { data: {} };
if (Array.isArray(data) && data.length >= 1) gpu = data[0];
let rows = `<tr><td style="color: var(--color-info)">GPU</td><td>${gpu.name || "unknown"}</td></tr>`;
for (const item of Object.entries(gpu.data)) rows += `<tr><td style="color: var(--color-info)">${item[0]}</td><td>${item[1]}</td></tr>`;
gpuTbody.innerHTML = rows;
if (gpu.chart && gpu.chart.length === 2) updateGPUChart(gpu.chart[0], gpu.chart[1]);
gpuEl.style.display = "block";
} catch (e) {
error("updateGPU", e);
@@ -19146,7 +19206,7 @@ async function startGPU() {
gpuEl.style.display = "block";
if (gpuInterval) clearInterval(gpuInterval);
const interval = window.opts?.gpu_monitor || 3e3;
log("startGPU", interval);
log("startGPUmonitor", interval);
gpuInterval = setInterval(updateGPU, interval);
updateGPU();
}
ui/dist/sdnext.mjs.map
Vendored
+2 -2
View File
File diff suppressed because one or more lines are too long
ui/extraNetworks.ts
+7
View File
@@ -4,12 +4,17 @@ import { authFetch } from './authWrap';
import { timer } from './timers';
const activePromptTextarea = {};
const selectedNetworks = {};
let sortVal = -1;
let totalCards = -1;
let lastTab = 'control';
// helpers
export function getSelectedNetworks() {
return selectedNetworks;
}
export const getENActiveTab = () => {
let tabName = '';
if (gradioApp().getElementById('txt2img_prompt')?.checkVisibility() || gradioApp().getElementById('txt2img_generate')?.checkVisibility()) tabName = 'txt2img';
@@ -307,6 +312,7 @@ function refreshENInput(tabName) {
export async function markSelectedCards(selected, page = '') {
log('markSelectedCards', selected, page);
selectedNetworks[page] = selected;
gradioApp().querySelectorAll('.extra-network-cards .card').forEach((el) => {
if (page.length > 0 && el.dataset.page !== page) return; // filter by page
if (selected.includes(el.dataset.name) || selected.includes(el.dataset.short)) el.classList.add('card-selected');
@@ -660,6 +666,7 @@ window.applyStyles = applyStyles;
window.closeDetailsEN = closeDetailsEN;
window.getENActivePage = getENActivePage;
window.getCardDetails = getCardDetails;
window.getSelectedNetworks = getSelectedNetworks;
window.sortExtraNetworks = sortExtraNetworks;
window.refeshDetailsEN = refeshDetailsEN;
window.extraNetworksSearchButton = extraNetworksSearchButton;
ui/gallery.ts
+52 -33
View File
@@ -34,7 +34,6 @@ const minCleanupCount = 1000;
const minCleanupTime = 1000 * 60 * 60; // 1 hour
const folderStylesheet = new CSSStyleSheet();
const fileStylesheet = new CSSStyleSheet();
const iconStopwatch = String.fromCodePoint(9201);
// Store separator states for the session
const separatorStates = new Map();
const el = {
@@ -49,6 +48,12 @@ const el = {
const cleanupTimers = {};
const maintenanceTimers = {};
const fetchQueue = [];
const icons = {
Time: String.fromCodePoint(9201),
Folder: String.fromCodePoint(128448), // or 128449;
Sort: String.fromCodePoint(8645),
Images: String.fromCodePoint(128461),
};
const SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'tiff', 'jp2', 'jxl', 'gif', 'mp4', 'mkv', 'avi', 'mjpeg', 'mpg', 'avr'];
@@ -63,6 +68,17 @@ const gallerySorter = {
modA: { name: 'Modified Descending', func: (b, a) => a.mtime - b.mtime },
none: { name: 'None', func: undefined },
};
const folderSorter = {
nameA: { name: 'Name Ascending', func: (a, b) => a.name.localeCompare(b.name) },
nameD: { name: 'Name Descending', func: (b, a) => a.name.localeCompare(b.name) },
sizeD: { name: 'Size Ascending', func: (a, b) => a.size - b.size },
sizeA: { name: 'Size Descending', func: (b, a) => a.size - b.size },
resD: { name: 'Resolution Ascending', func: (a, b) => a.width * a.height - b.width * b.height },
resA: { name: 'Resolution Descending', func: (b, a) => a.width * a.height - b.width * b.height },
modD: { name: 'Modified Ascending', func: (a, b) => a.mtime - b.mtime },
modA: { name: 'Modified Descending', func: (b, a) => a.mtime - b.mtime },
none: { name: 'None', func: undefined },
};
let sortMode = gallerySorter.none;
@@ -225,9 +241,9 @@ function updateGalleryStyles() {
padding: 4px;
font-size: 1.2em;
letter-spacing: 0.5em;
width: 140px;
margin-top: calc(140px - 32px);
opacity: 75%;
border-radius: var(--sd-border-radius);
}
:host(.gallery-file-selected) .gallery-file {
box-shadow: 0 0 0 2px var(--sd-button-selected-color);
@@ -278,7 +294,6 @@ class SimpleProgressBar {
this.#progress.style.backgroundColor = 'var(--sd-main-accent-color)';
this.#textDiv.style.cssText = 'position:relative; margin:auto; width:max-content; height:100%;';
this.#text.style.cssText = 'user-select:none; color:white;';
this.#textDiv.append(this.#text);
this.#container.append(this.#progress, this.#textDiv);
}
@@ -287,6 +302,10 @@ class SimpleProgressBar {
if (total <= 0) return;
this.hide();
this.#max = total;
if (this.#monitoredSet.size >= this.#max) {
this.stop();
return;
}
this.#interval = setInterval(() => this.update(this.#monitoredSet.size, this.#max), 100);
}
@@ -314,7 +333,7 @@ class SimpleProgressBar {
stop() {
clearInterval(this.#interval);
this.#interval = undefined;
// this.#interval = undefined;
if (this.stats.count) {
debug('gallery: thumbnail stats', this.stats);
this.stats = { ...this.defaultStats };
@@ -629,7 +648,7 @@ class GalleryFile extends HTMLElement {
img.title = `Folder: ${this.folder}\nFile: ${this.name}\nSize: ${this.size.toLocaleString()} bytes\nModified: ${this.mtime.toLocaleString()}`;
this.title = img.title;
const shouldDisplayBasedOnSearch = this.title.toLowerCase().includes(el.search.value.toLowerCase()); // Final visibility check based on search term.
if (this.style.display !== 'none') this.style.display = shouldDisplayBasedOnSearch ? 'flex' : 'none'; // Only proceed if not already hidden by a closed separator
if (this.style.display !== 'none') this.style.display = shouldDisplayBasedOnSearch ? '' : 'none'; // Only proceed if not already hidden by a closed separator
this.shadow.appendChild(img);
pb.stats.elapsed = (pb.stats.elapsed || 0) + Math.round(performance.now() - t0);
}
@@ -674,7 +693,7 @@ async function handleSeparator(separator) {
const fileDirPath = fileDir ? fileDir[1] : '';
if (separator.title.length > 0 && fileDirPath === separator.title) {
f.style.display = nowHidden ? 'none' : 'unset';
f.style.display = nowHidden ? 'none' : '';
}
}
// Note: Count is not updated here on manual toggle, as it reflects the total.
@@ -774,24 +793,22 @@ window.gallerySendImage = gallerySendImage;
* @param {...string|[string, string]} messages - Each can be either a string to use as-is, or an array of a string label and value
* @returns {void}
*/
function updateStatusWithSort(...messages) {
function updateStatusLine(...messages) {
if (!el.status) return;
messages.unshift(['Sort', sortMode.name]);
const fragment = document.createDocumentFragment();
for (let i = 0; i < messages.length; i++) {
const div = document.createElement('div');
if (Array.isArray(messages[i])) {
const [text1, text2] = messages[i];
const tDiv1 = document.createElement('div');
tDiv1.innerText = `${text1}:`;
const tDiv2 = document.createElement('div');
tDiv2.innerText = text2;
tDiv2.title = text2;
div.append(tDiv1, tDiv2);
const [k, v] = messages[i];
const tDiv = document.createElement('div');
const ico = icons[k] || `${k}:`;
tDiv.innerText = `${ico} ${v}`;
div.append(tDiv);
} else {
const tDiv1 = document.createElement('div');
tDiv1.innerText = messages[i];
div.append(tDiv1);
const tDiv = document.createElement('div');
tDiv.innerText = messages[i];
div.append(tDiv);
}
fragment.append(div);
}
@@ -878,7 +895,7 @@ async function gallerySearch() {
f.style.display = (!dirPath || isOpen) ? 'unset' : 'none';
});
updateStatusWithSort('Filter', 'Cleared', ['Images', allFiles.length.toLocaleString()]);
updateStatusLine('Filter', 'Cleared', ['Images', allFiles.length.toLocaleString()]);
return;
}
@@ -936,7 +953,7 @@ async function gallerySearch() {
}
const t1 = performance.now();
updateStatusWithSort('Filter', ['Images', `${totalFound.toLocaleString()} / ${allFiles.length.toLocaleString()}`], `${iconStopwatch} ${Math.round(t1 - t0).toLocaleString()}ms`);
updateStatusLine('Filter', ['Images', `${totalFound.toLocaleString()} / ${allFiles.length.toLocaleString()}`], ['Time', `${Math.round(t1 - t0).toLocaleString()}ms`]);
timer(`galleryFilter:${str}`, t1 - t0);
refreshGallerySelection();
}, 250);
@@ -981,9 +998,7 @@ export async function gallerySort(key) {
const folderGroups = new Map();
for (const file of subfolderFiles) {
const dir = getDirPath(file);
if (!folderGroups.has(dir)) {
folderGroups.set(dir, []);
}
if (!folderGroups.has(dir)) folderGroups.set(dir, []);
folderGroups.get(dir).push(file);
}
@@ -994,7 +1009,11 @@ export async function gallerySort(key) {
rootFiles.forEach((node) => fragment.appendChild(node));
// Sort folder names alphabetically, then sort files within each folder
const sortedFolderNames = Array.from<any>(folderGroups.keys()).sort((a, b) => a.localeCompare(b));
// const sortedFolderNames = Array.from<any>(folderGroups.keys()).sort((a, b) => a.localeCompare(b));
const folderNames = Array.from<string>(folderGroups.keys());
const sortedFolderNames = currentSort.endsWith('A') ? folderNames.sort((a, b) => a.localeCompare(b)) : folderNames.sort((a, b) => b.localeCompare(a));
console.log('HERE', sortedFolderNames);
for (const folderName of sortedFolderNames) {
const files = folderGroups.get(folderName);
files.sort(sortMode.func);
@@ -1023,7 +1042,7 @@ export async function gallerySort(key) {
const t1 = performance.now();
log(`gallerySort: sort=${sortMode.name} len=${arr.length} time=${Math.floor(t1 - t0)}`);
updateStatusWithSort(['Images', arr.length.toLocaleString()], `${iconStopwatch} ${Math.round(t1 - t0).toLocaleString()}ms`);
updateStatusLine(['Images', arr.length.toLocaleString()], ['Time', `${Math.round(t1 - t0).toLocaleString()}ms`]);
timer(`gallerySort:${sortMode.name}`, t1 - t0);
refreshGallerySelection();
}
@@ -1120,7 +1139,7 @@ async function thumbCacheCleanup(folder, imgCount, controller, force = false) {
log('galleryMaintenance', { folder, kept: keptGalleryHashes.size, deleted: delcount, time: Math.round(t1 - t0) });
timer(`thumbnailDBCleanup:${folder}`, t1 - t0);
currentGalleryFolder = null;
updateStatusWithSort('Thumbnail cache cleared');
updateStatusLine('Thumbnail cache cleared');
})
.catch((reason) => {
SimpleFunctionQueue.abortLogger('thumbCacheCleanup', reason);
@@ -1157,7 +1176,7 @@ function clearCacheIfDisabled(browser_cache) {
.then(() => {
log('thumbCacheCleanup', { time: Math.floor(performance.now() - t0) });
currentGalleryFolder = null;
updateStatusWithSort('Thumbnail cache cleared');
updateStatusLine('Thumbnail cache cleared');
})
.catch((e) => {
SimpleFunctionQueue.abortLogger('thumbCacheCleanup', e);
@@ -1183,12 +1202,12 @@ window.clearCache = clearCache;
async function fetchFilesHT(evt, controller) {
const t0 = performance.now();
const fragment = document.createDocumentFragment();
updateStatusWithSort(['Folder', evt.target.name], 'in-progress');
updateStatusLine(['Folder', evt.target.name], 'in-progress');
let numFiles = 0;
const res = await authFetch(`${window.api}/browser/files?folder=${encodeURI(evt.target.name)}`);
if (!res || res.status !== 200) {
updateStatusWithSort(['Folder', evt.target.name], ['Failed', res?.statusText || 'No response']);
updateStatusLine(['Folder', evt.target.name], ['Failed', res?.statusText || 'No response']);
return;
}
const jsonData = await res.json();
@@ -1209,7 +1228,7 @@ async function fetchFilesHT(evt, controller) {
const t1 = performance.now();
log(`gallery: folder=${evt.target.name} num=${numFiles} method=http time=${Math.floor(t1 - t0)}ms`);
timer(`galleryFetch:${evt.target.name}`, t1 - t0);
updateStatusWithSort(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], ['Time', `${Math.floor(t1 - t0).toLocaleString()}ms`]);
pb.start(numFiles);
addSeparators();
refreshGallerySelection();
@@ -1238,7 +1257,7 @@ async function fetchFilesWS(evt) { // fetch file-by-file list over websockets
await fetchFilesHT(evt, controller); // fallback to http
return;
}
updateStatusWithSort(['Folder', evt.target.name]);
updateStatusLine(['Folder', evt.target.name]);
const t0 = performance.now();
let numFiles = 0;
let t1 = performance.now();
@@ -1257,7 +1276,7 @@ async function fetchFilesWS(evt) { // fetch file-by-file list over websockets
numFiles++;
fragment.appendChild(file);
if (numFiles % fragmentSize === 0) {
updateStatusWithSort(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], 'in-progress', `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], ['Status', 'in-progress'], ['Time', `${Math.floor(t1 - t0).toLocaleString()}ms`]);
el.files.appendChild(fragment);
fragment = document.createDocumentFragment();
}
@@ -1269,7 +1288,7 @@ async function fetchFilesWS(evt) { // fetch file-by-file list over websockets
el.files.appendChild(fragment);
// gallerySort();
log(`gallery: folder=${evt.target.name} num=${numFiles} method=ws time=${Math.floor(t1 - t0)}ms`);
updateStatusWithSort(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], `${iconStopwatch} ${Math.floor(t1 - t0).toLocaleString()}ms`);
updateStatusLine(['Folder', evt.target.name], ['Images', numFiles.toLocaleString()], ['Time', `${Math.floor(t1 - t0).toLocaleString()}ms`]);
pb.start(numFiles);
addSeparators();
refreshGallerySelection();
@@ -1362,7 +1381,7 @@ async function initGalleryAutoRefresh() {
}
async function overlayDelete(evt) {
const res = await authFetch(`${window.api}/delete-image?file=${encodeURIComponent(currentImage)}`);
const res = await authFetch(`${window.api}/delete-image?file=${encodeURIComponent(currentImage)}`, { method: 'DELETE' });
evt.stopPropagation();
if (!res || res.status !== 200) {
error('galleryDelete', { file: currentImage, status: res?.status, statusText: res?.statusText });
ui/globals.d.ts
Vendored
+1
View File
@@ -102,6 +102,7 @@ declare global {
closeDetailsEN?: (...args: unknown[]) => unknown[]; // ui/extraNetworks.ts
extraNetworksFilterVersion?: (event: Event) => void; // ui/extraNetworks.ts
extraNetworksSearchButton?: (event: Event) => void; // ui/extraNetworks.ts
getSelectedNetworks?: () => Record<string, string[]>; // ui/extraNetworks.ts
getCardDetails?: (...args: unknown[]) => unknown[]; // ui/extraNetworks.ts
quickSaveStyle?: () => void; // ui/extraNetworks.ts
getENActivePage?: () => string; // ui/extraNetworks.ts

Some files were not shown because too many files have changed in this diff Show More