Merge pull request #5002 from vladmandic/dev

merge dev

CHANGELOG.md

@@ -1,10 +1,50 @@
 # Change Log for SD.Next
 
+## Update for 2026-07-23
+
+Primarily a service release with updates to compute packages: torch, CUDA, ROCm, etc.
+Plus optimizations to SDNQ quantization and attention  
+And update to process tab, several quality-of-life improvements and bug-fixes  
+
+[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)  
+
+- **Compute**
+  - torch: update to `2.13.0` for CUDA, ROCm, IPEX
+  - torch: explicitly set inductor and triton cache locations
+  - torch: log triton/dynamo/inductor timer stats
+  - cuda: update to `13.2`
+  - sdnq quantization optimizations
+  - sdnq attention optimizations
+  - sdnq separate dit/te settings
+- **Features**
+  - process: read video properties and metadata
+  - process: allow processing of video files  
+    *note*: currently only seedvr postprocessing is supported  
+    other workflows will be added in future releases  
+  - seedvr: enhanced upscaler support
+  - logs: propagate server tracebacks to client
+  - networks: improve search and filtering to allow multi-words  
+  - hotkeys: add alt+0-9 to switch to tab 0-9
+- **Fixes**
+  - attention: skip reapply
+  - download: better matching of shared components
+  - gallery: send to caption
+  - hotkeys: legacy-vs-modernui
+  - kanvas: paint combined with zoom
+  - load: flux1 t5
+  - logger: handle invalid subsystem log messages
+  - lora: support diffusers trainer
+  - preview: acknowledge visible/hidden on finish
+  - preview: cache image for reuse
+  - process: generate button busy tracking
+  - rembg: numba dependencies
+  - upscaler: auto-refresh to catch chainner upscalers that are not loaded on first attempt  
+
 ## 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:  
+*What's New?* Full week(!) since the last 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

TODO.md

@@ -13,16 +13,18 @@
 - Cloud providers, @CalamitousFelicitousness
 - Video processing add/verify full API support, @CalamitousFelicitousness
 - Storage analyzer, @vladmandic
+- Lora: new handler, @CalamitousFelicitousness
+- Vide: full prompt enhance
+- Processing -> Video capabilities, @vladmandic
+  - `NudeNet` in processing
+  - `RIFE` in processing
 
 ### Unassigned
 
 - [Nunchaku Lite](https://github.com/huggingface/diffusers/pull/14100)
-- Processing -> Video capabilities
-  - `RIFE` in processing
-  - `SeedVR2` in processing
 - [Object clear](https://huggingface.co/jixin0101/ObjectClear) remover for Kanvas
 - Video models: add to Reference
-- Video models: support custom entries
+- Video models: support custom entries, finetunes
 - UI Lite vs Expert mode
 - Auto handle scheduler `prediction_type`
 - Cache models in memory

data/previews.json

@@ -32,8 +32,10 @@
   "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--Flux.2-Klein-9B-KV-Merge-sdnq-hadamard-uint4": "black-forest-labs--FLUX.2-klein-9b-kv.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--Krea-2-Turbo-Merge-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",

installer.py

@@ -638,7 +638,7 @@ def install_cuda():
     if args.use_nightly:
         cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu132 --extra-index-url https://download.pytorch.org/whl/nightly/cu130')
     else:
-        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.12.0+cu130 torchvision==0.27.0+cu130 --index-url https://download.pytorch.org/whl/cu130')
+        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+cu132 torchvision==0.28.0+cu132 --index-url https://download.pytorch.org/whl/cu132')
     return cmd
 
 
@@ -720,7 +720,8 @@ def install_rocm_zluda():
                 check_python(supported_minors=[12], reason='ROCm-Windows: preview python==3.12 required')
                 # torch 2.8.0a0 is the last version with rocm 6.4 support
                 torch_command = os.environ.get('TORCH_COMMAND', '--no-cache-dir https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torch-2.8.0a0%2Bgitfc14c65-cp312-cp312-win_amd64.whl https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torchvision-0.24.0a0%2Bc85f008-cp312-cp312-win_amd64.whl')
-    else:
+
+    else: # linux
         #check_python(supported_minors=[10, 11, 12, 13, 14], reason='ROCm backend requires a Python version between 3.10 and 3.13')
         if args.use_nightly:
             if rocm.version is None or float(rocm.version) >= 7.2: # assume the latest if version check fails
@@ -729,9 +730,9 @@ def install_rocm_zluda():
                 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm7.1')
         else:
             if rocm.version is None or float(rocm.version) >= 7.2: # assume the latest if version check fails
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.12.0+rocm7.2 torchvision==0.27.0+rocm7.2 --index-url https://download.pytorch.org/whl/rocm7.2')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+rocm7.2 torchvision==0.28.0+rocm7.2 --index-url https://download.pytorch.org/whl/rocm7.2')
             elif rocm.version == "7.1":
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.12.0+rocm7.1 torchvision==0.27.0+rocm7.1 --index-url https://download.pytorch.org/whl/rocm7.1')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+rocm7.1 torchvision==0.28.0+rocm7.1 --index-url https://download.pytorch.org/whl/rocm7.1')
             elif rocm.version == "7.0":
                 torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.10.0+rocm7.0 torchvision==0.25.0+rocm7.0 --index-url https://download.pytorch.org/whl/rocm7.0')
             elif rocm.version == "6.4":
@@ -769,7 +770,7 @@ def install_ipex():
     if args.use_nightly:
         torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/xpu')
     else:
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.12.0+xpu torchvision==0.27.0+xpu --extra-index-url https://download.pytorch.org/whl/xpu')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+xpu torchvision==0.28.0+xpu --extra-index-url https://download.pytorch.org/whl/xpu')
 
     ts('ipex', t_start)
     return torch_command
@@ -785,7 +786,7 @@ def install_openvino():
         torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.11.0+cpu torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cpu')
 
     if not (args.skip_all or args.skip_requirements):
-        install(os.environ.get('OPENVINO_COMMAND', 'openvino==2026.1.0'), 'openvino')
+        install(os.environ.get('OPENVINO_COMMAND', 'openvino==2026.2.1'), 'openvino')
     ts('openvino', t_start)
     return torch_command
 

modules/api/caption.py

@@ -678,6 +678,7 @@ def get_vqa_prompts(model: str | None = None):
     - **promptgen**: Analyze, Generate Tags, Mixed Caption, Mixed Caption+ (MiaoshouAI PromptGen fine-tunes only)
     - **moondream**: Point at..., Detect all... (Moondream 2 and 3)
     - **moondream2_only**: Detect Gaze (Moondream 2 only)
+    - **toriigate**: Long Thoughts, Structured Markdown, JSON Caption, Comic Markdown, Chroma Style and other native caption formats (ToriiGate 0.5 only)
     """
     from modules.caption import vqa
     if model:
@@ -688,7 +689,8 @@ def get_vqa_prompts(model: str | None = None):
         "florence": vqa.vlm_prompts_florence,
         "promptgen": vqa.vlm_prompts_promptgen,
         "moondream": vqa.vlm_prompts_moondream,
-        "moondream2_only": vqa.vlm_prompts_moondream2
+        "moondream2_only": vqa.vlm_prompts_moondream2,
+        "toriigate": vqa.vlm_prompts_toriigate
     }
 
 

modules/api/process.py

@@ -292,7 +292,7 @@ class APIProcess:
         reqDict, script_args = self.set_upscalers(req)
         reqDict['image'] = helpers.decode_base64_to_image(reqDict['image'])
         with self.queue_lock:
-            result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, script_args=script_args, **reqDict)
+            result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", video="", save_output=False, script_args=script_args, **reqDict)
         return models.ResProcessImage(image=helpers.encode_pil_to_base64(result[0][0]), html_info=result[1])
 
     def extras_batch_images_api(self, req: models.ReqProcessBatch):
@@ -301,5 +301,5 @@ class APIProcess:
         image_list = reqDict.pop('imageList', [])
         image_folder = [helpers.decode_base64_to_image(x.data) for x in image_list]
         with self.queue_lock:
-            result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, script_args=script_args, **reqDict)
+            result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", video="", save_output=False, script_args=script_args, **reqDict)
         return models.ResProcessBatch(images=list(map(helpers.encode_pil_to_base64, result[0])), html_info=result[1])

modules/attention.py

@@ -1,6 +1,6 @@
 from functools import wraps
 import torch
-from modules import rocm, errors
+from modules import rocm, errors, devices
 from modules.logger import log
 from installer import install, installed, torch_info
 
@@ -26,7 +26,8 @@ def set_sdnq_attention():
         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[-2] >= 512 or key.shape[-2] >= 512) # Skip TE
+                and (query.shape[-2] >= 32 and key.shape[-2] >= 32)
+                and (query.shape[-2] > 512 or key.shape[-2] > 512) # Skip TE
                 and query.shape[-3] > 1 # Skip VAE
             ):
                 return sdnq_triton_atten(
@@ -37,7 +38,6 @@ def set_sdnq_attention():
                     smooth_k=shared.opts.sdnq_attention_smooth_k,
                     use_hadamard=shared.opts.sdnq_attention_use_hadamard,
                     hadamard_group_size=shared.opts.sdnq_attention_hadamard_group_size,
-                    do_quantize=shared.opts.sdnq_attention_use_quantized_matmul,
                 )
             else:
                 if enable_gqa:
@@ -45,7 +45,7 @@ def set_sdnq_attention():
                 return sdpa_pre_sdnq_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
         torch.nn.functional.scaled_dot_product_attention = sdpa_sdnq_atten
         torch_info.set(attention='sdnq')
-        log.debug('Torch attention: type="SDNQ attention"')
+        log.debug(f'Torch attention: type="SDNQ attention" matmul={shared.opts.sdnq_attention_matmul_type}:{shared.opts.sdnq_attention_pv_matmul_type} smooth={shared.opts.sdnq_attention_smooth_k} hadamard={shared.opts.sdnq_attention_use_hadamard}')
     except Exception as err:
         log.error(f'Torch attention: type="SDNQ attention" {err}')
 
@@ -262,7 +262,7 @@ def set_diffusers_attention(pipe, quiet = False):
     if shared.opts.cross_attention_optimization == "Disabled":
         torch_info.set(attention="disabled")
     elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product":  # The default set by Diffusers
-        torch_info.set(attention="sdpa")
+        devices.set_sdpa_params()
         # set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product")
     elif shared.opts.cross_attention_optimization == "xFormers":
         if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):

modules/caption/models_def.py

@@ -1,4 +1,5 @@
 from modules import ui_symbols
+from modules.caption.toriigate import prompt_list as vlm_prompts_toriigate
 
 
 vlm_models = {
@@ -17,6 +18,7 @@ vlm_models = {
     "Alibaba Qwen 3.5 4B": "Qwen/Qwen3.5-4B",
     "Alibaba Qwen 3.5 2B": "Qwen/Qwen3.5-2B",
     "Alibaba Qwen 3.5 0.8B": "Qwen/Qwen3.5-0.8B",
+    "ToriiGate 0.5": "Minthy/ToriiGate-0.5",
     "JoyTag": "fancyfeast/joytag",
     "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava",
     "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava",
@@ -173,10 +175,11 @@ vlm_prompt_placeholders = {
     "Point at...": "Enter objects to locate, e.g., 'the red car' or 'all the eyes'",
     "Detect all...": "Enter object type to detect, e.g., 'cars' or 'faces'",
     "Detect Gaze": "No input needed - auto-detects face and gaze direction",
+    **{task: "No input needed - ToriiGate uses the selected caption format" for task in vlm_prompts_toriigate},
 }
 
 # Legacy list for backwards compatibility
-vlm_prompts = vlm_prompts_common + vlm_prompts_florence + vlm_prompts_promptgen + vlm_prompts_moondream + vlm_prompts_moondream2
+vlm_prompts = vlm_prompts_common + vlm_prompts_florence + vlm_prompts_promptgen + vlm_prompts_moondream + vlm_prompts_moondream2 + vlm_prompts_toriigate
 
 vlm_prefill = 'Answer: the image shows'
 

modules/caption/toriigate.py

@@ -0,0 +1,208 @@
+# Prompt definitions for ToriiGate 0.5 (Minthy/ToriiGate-0.5)
+# The model is a Qwen3.5 vision fine-tune trained on one exact query structure and one system prompt:
+#   system: captioning expert
+#   user:   "# Captioning format:\n<template>\n\n# Characters on picture:\n<names instruction>"
+# Templates are reproduced verbatim from the model repo (scripts/prompts.py); the model degrades
+# when they are paraphrased. Loading and generation stay on the shared Qwen path in vqa.py.
+
+system_prompt = "You are image captioning expert. Describe user's picture according to requested format and instructions."
+
+names_instruction = "Try to recognize the characters in the picture and use their names."
+no_names_instruction = "Avoid to guess names for characters."
+
+# format key -> template, verbatim from the model repo
+formats = {
+"long_thoughts_v2": """Your answer must contain 6 parts:
+<format>
+# 1. Thoughts about characters
+You need to think here and compare peoples/creatures that you see on the picture with given popular tags, or descriptions, or your memories for each characters to determine who is who.
+# 2. Key details
+Here you need to determine key details on comic and list them.
+# 3. Long description
+Here come up with a long and detailed description of image content. Be creative, mention all detailes you listed above and other important things.
+# 4. Detailed description for each character
+## Name 1
+Detailed and long description for the first character
+## Name 2
+Same for each one (if present)
+</format>
+""",
+"long_thoughts": """Your answer must contain 6 parts:
+<format>
+# 1. Thoughts about characters
+You need to think here and compare peoples/creatures that you see on the picture  with given popular tags, or descriptions, or your memories for each characters to determine who is who.
+If no characters are listed in input - just write here "No named characters"
+# 2. General description
+A one-two paragraph summary of the image. Mention all individual parts/objects/characters/positions/interactions/etc.
+# 3. Detailed description for each character
+## Character name 1 (put here the name if any)
+In very detail write about features, poses, look, used objects, interactions, and other things for character on the picture.
+## Character name 2 (put here the name if any)
+Same for each character.
+...
+# 4. Individual Parts
+List the individual things you see in the image and their relative positions to other parts. Use a numbered list of between 5 and 20 items depending on image complexity.
+# 5. Texts on image
+Mention every texts that you notice on image, including types (a speech bubble, watermark, banner, etc.) and content.
+# 6. Background and effects
+Give some info about objects on background, describe the location (if seen). Then mention effects (style, camera angle, clarity/blurrines, effects like depth of field, strange angle/forshortening, etc.)
+</format>
+""",
+"json": """Use json-style caption for given image with following structure:
+{"character" : "Description for character or object. Name (if defined), main details, features, position, pose, etc.",
+/or in case of multiple
+"character_1" : "Description for first"
+"character_2" : "Description for second ",
+"character_N"...
+/or if there are no characters
+"main content" : "long and detailed description of main content of image that might be the main focus if characters are missing",
+/
+"background" : "Detailed descritpion of background and it's content",
+"image_effects" : "If there are some visual effects like fisheye distortion, chromatic aberration, glitches, messy drawing or anything else - write about it. If it's just a general anime art - omit this field."
+"texts" : "Speech bubbles, bars, marks, signs etc. with texts if present, else None",
+"atmosphere" : "...",
+}
+In special cases you can add extra keys.
+""",
+"long": """Make a caption for given image with natural text. Use 2 to 5 paragraphs. Make your description long and vivid, mentioning all the details.
+""",
+"min_structured_md": """Your answer must contain 3 parts:
+<format>
+# 1. Thoughts about characters
+You need to think here and compare peoples/creatures that you see on the picture  with given popular tags, or descriptions, or your memories for each characters to determine who is who.
+If no characters are listed in input - just write here "No named characters"
+# 2. Key details
+Here you need to write about the key details on image, prefere using regular text.
+# 3. Structured description
+## General
+Write about general composition, content of image, background and all things that are not related to characters directly.
+## Character name 1 (put here the name if any)
+Write about datails and content related to specific character, including features, poses, look, used objects, interactions, and other things.
+## Character name 2 (put here the name if any)
+Same for each character.
+## Image effects
+Mention image effect, style, camera angle
+</format>
+In general stick to shorter descriptions.
+""",
+"json_comic": """Use json-style caption to describe to comin, stick to following structure:
+{
+"comic_format": "menation the format, for example Comic of N frames",
+"1st_frame": "Main description of the content for fist frame",
+"2nd_frame": "Same for the second",
+...
+"Nth_ftame": "...",
+"character_1": "Describe the characters in comic",
+...
+"character_N": "Separate description for each",
+"meaning": "Try to guess general mood, vibe and meaning of the comic"
+}
+""",
+"md_comic": """Use markdown format to describe to comic, 5 parts are recommended:
+<format>
+# 1. Thoughts about characters
+You need to think here and compare peoples/creatures that you see on the picture with given popular tags, or descriptions, or your memories for each characters to determine who is who.
+# 2. Key details
+Here you need to determine key details on comic and list them.
+# 3. Comic format
+In this section come up with the description of comic format, how many pages there are, horisontal/vertical orientation and other things. Optionally you can list main characters here.
+# 4. Details for each frame
+## 4.1 Frame 1 (position)
+Description for each frame, includding characters, objects, interactions, texts/speech bubbles and other things. Be detailed but not overdoo.
+## 4.2 Frame 2 (position)
+Same for each frame.
+...
+# 5. Extra comment
+Here you should write general desciption and some other info about the image.
+</format>
+""",
+"min_structured_json": """
+Use json-style caption for given image with following structure:
+{"General" : "Here you need to come up with general/common information about picture, overall composition. Stick to shorter phrases and tags instead of long purple prose. Avoid bullets and markdown, write in plain text.",
+"character_1 (put here the name if any)" : "Description of first character."
+"character_2 (if present" : "Description for second ",
+"character_N"
+...
+"image_effects" : "Mention here effects on image if there are any distinct."
+"texts" : "Speech bubbles, bars, marks, signs etc. with texts if present, else None",
+"watermarks" : "If present",
+}
+Prefere shorter description and tags.
+""",
+"chroma-style": """Your task is to describe the picture in very detail using a structure of 4 parts.
+### 1. Regular Summary:
+[A one-paragraph summary of the image. The paragraph should mention all individual parts/things/characters/etc.]
+### 2. Individual Parts:
+[List the individual things you see in the image and their relative positions to other parts. Use a numbered list of between 5 and 30 items depending on image complexity.]
+### 3. Midjourney-Style Summary:
+[A summary that has higher concept density by using comma-separated partial sentences instead of proper sentence structure.]
+### 4. DeviantArt Commission Request
+[Write a description as if you're commissioning this *exact* image via someone who is currently taking requests.]
+""",
+"short": """The caption for image should be quite short without long purple prose and slop. Cover main objects and details.
+""",
+}
+
+# task label shown in the UI -> format key; the first entry is what the task dropdown falls back to
+tasks = {
+    "Long Thoughts": "long_thoughts_v2",
+    "Long Thoughts Full": "long_thoughts",
+    "Structured Markdown": "min_structured_md",
+    "Structured JSON": "min_structured_json",
+    "JSON Caption": "json",
+    "Comic Markdown": "md_comic",
+    "Comic JSON": "json_comic",
+    "Chroma Style": "chroma-style",
+}
+
+# common tasks reach the handler as internal tokens; Normal Caption has no ToriiGate format and is not offered,
+# but it stays mapped here because the API accepts any task for any model
+common_formats = {
+    "<CAPTION>": "short",
+    "<DETAILED_CAPTION>": "long",
+    "<MORE_DETAILED_CAPTION>": "long",
+    "Short Caption": "short",
+    "Normal Caption": "long",
+    "Long Caption": "long",
+}
+
+# formats whose reasoning block only produces results when character names are requested
+names_only = {"long_thoughts_v2", "long_thoughts", "md_comic", "min_structured_md"}
+
+prompt_list = list(tasks)
+
+
+def is_toriigate(name: str) -> bool:
+    """Match ToriiGate 0.5 by display name or repo id; the 0.4 fine-tunes use a different prompt format."""
+    if not name:
+        return False
+    return 'toriigate 0.5' in name.lower().replace('-', ' ')
+
+
+def resolve_format(question: str) -> tuple[str, str]:
+    """Map an incoming question to a format key, or to free text used as the format block.
+
+    Returns (format_key, custom_text). A non-empty custom_text replaces the stored template, so a
+    free-text question is answered in the requested form rather than as a generic caption.
+    """
+    question = (question or '').strip()
+    if not question:
+        return next(iter(tasks.values())), ''
+    if question in tasks:
+        return tasks[question], ''
+    if question in common_formats:
+        return common_formats[question], ''
+    return '', question
+
+
+def build_query(question: str, use_names: bool = True) -> tuple[str, str]:
+    """Build the (system, user) pair the model was trained on."""
+    fmt, custom = resolve_format(question)
+    template = custom or formats.get(fmt, formats['long'])
+    if fmt in names_only:
+        use_names = True  # the reasoning block of these formats returns nothing without names
+    query = '# Captioning format:\n'
+    query += template.rstrip('\n') + '\n\n'  # free text has no trailing newline of its own
+    query += '# Characters on picture:\n'
+    query += f'{names_instruction if use_names else no_names_instruction}\n'
+    return system_prompt, query

modules/caption/vqa.py

@@ -12,9 +12,9 @@ from PIL import Image
 from modules import shared, devices, errors, model_quant, sd_models, sd_models_compile
 from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
 from modules.logger import log, console
-from modules.caption import vqa_detection, helpers
+from modules.caption import vqa_detection, helpers, toriigate
 from modules.caption.attention import set_attention
-from modules.caption.models_def import vlm_models, vlm_system, vlm_analyze, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_reverse_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, analyze_question, get_vlm_repo # pylint: disable=unused-import
+from modules.caption.models_def import vlm_models, vlm_system, vlm_analyze, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_reverse_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, vlm_prompts_toriigate, analyze_question, get_vlm_repo # pylint: disable=unused-import
 
 
 debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
@@ -73,6 +73,10 @@ def get_prompts_for_model(model_name: str) -> list:
     if 'florence' in model_lower:
         return vlm_prompts_common + vlm_prompts_florence
 
+    # Check for ToriiGate 0.5 (native caption formats first, Normal Caption has no equivalent format)
+    if toriigate.is_toriigate(model_name):
+        return vlm_prompts_toriigate + [p for p in vlm_prompts_common if p != 'Normal Caption']
+
     # Check for Moondream models (Moondream 2 has gaze detection, Moondream 3 does not)
     if 'moondream' in model_lower:
         if 'moondream3' in model_lower or 'moondream 3' in model_lower:
@@ -140,6 +144,16 @@ def is_thinking_model(model_name: str) -> bool:
     return any(indicator in model_lower for indicator in thinking_indicators)
 
 
+def check_linear_attention(model):
+    """Warn when a hybrid linear-attention model lacks its kernels and falls back to a per-token torch loop."""
+    model_type = getattr(getattr(model, 'config', None), 'model_type', '') or ''
+    if not model_type.startswith('qwen3_5'):
+        return
+    from transformers.utils.import_utils import is_flash_linear_attention_available
+    if not is_flash_linear_attention_available():
+        log.warning(f'LLM: cls={model.__class__.__name__} linear attention kernels missing: install="flash-linear-attention" impact="generation runs a slow torch fallback"')
+
+
 def truncate_b64_in_conversation(conversation, front_chars=50, tail_chars=50, threshold=200):
     """
     Deep copy a conversation structure and truncate long base64 image strings for logging.
@@ -561,6 +575,7 @@ class VQA:
                 self.model = sd_models_compile.compile_torch(self.model, apply_to_components=False, op="VQA")
             register_aux('vqa', self.model)
             set_attention(self.model)
+            check_linear_attention(self.model)
             self.loaded = repo
             devices.torch_gc()
 
@@ -571,13 +586,18 @@ class VQA:
         cls_name = self.model.__class__.__name__
         debug(f'LLM: handler=qwen model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}')
 
-        question = question.replace('<', '').replace('>', '').replace('_', ' ')
-        if question is not None and len(question) > 4:
-            if question in vlm_prompt_reverse_mapping:
-                debug(f'LLM: handler=gemma mapping friendly question="{question}" to internal="{vlm_prompt_reverse_mapping[question]}"')
-                question = vlm_prompt_reverse_mapping[question]
-
         system_prompt = system_prompt or shared.opts.caption_vlm_system
+        if toriigate.is_toriigate(repo):
+            # ToriiGate 0.5 is trained on one system prompt and one query structure and degrades on anything else,
+            # so its own prompts replace both the user system prompt and the generic token cleanup
+            system_prompt, question = toriigate.build_query(question)
+            debug(f'LLM: handler=qwen toriigate system="{system_prompt}" query="{question}"')
+        else:
+            question = question.replace('<', '').replace('>', '').replace('_', ' ')
+            if question is not None and len(question) > 4:
+                if question in vlm_prompt_reverse_mapping:
+                    debug(f'LLM: handler=gemma mapping friendly question="{question}" to internal="{vlm_prompt_reverse_mapping[question]}"')
+                    question = vlm_prompt_reverse_mapping[question]
         conversation = [
             {
                 "role": "system",

modules/control/util.py

@@ -161,6 +161,8 @@ def blend(images):
 
 
 def decode_fourcc(cc):
+    if cc is None:
+        return None
     cc_bytes = int(cc).to_bytes(4, byteorder=sys.byteorder) # convert code to a bytearray
     cc_str = cc_bytes.decode() # decode byteaarray to a string
     return cc_str

modules/detailer/detailer.py

@@ -307,6 +307,8 @@ class Detailer():
         via detailer_opt(). The seed is resolved here so restore()'s inpaint passes are reproducible and the
         effective value can be reported back.
         """
+        if image is None:
+            return None
         from modules.processing_helpers import get_fixed_seed
         from modules.paths import resolve_output_path
         seed = int(get_fixed_seed(seed))

modules/devices.py

@@ -432,7 +432,6 @@ def test_triton(early: bool = False):
         if triton_version is None:
             try:
                 import torch._inductor.triton as torch_triton
-
                 triton_version = torch_triton.__version__
             except Exception:
                 pass
@@ -542,10 +541,11 @@ def set_sdpa_params():
             sage = version('sageattention')
         except Exception:
             sage = False
-        log.debug(f'Torch attention installed: flashattn={flash} sageattention={sage}')
+        if flash or sage:
+            log.debug(f'Torch attention installed: flashattn={flash} sageattention={sage}')
 
         from diffusers.models import attention_dispatch as a
-        log.debug(f'Torch attention status: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} aiter={a._CAN_USE_AITER_ATTN} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN} kernels={a.is_kernels_available()}') # pylint: disable=protected-access
+        log.debug(f'Torch attention available: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} aiter={a._CAN_USE_AITER_ATTN} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN} kernels={a.is_kernels_available()} sdnq=True') # pylint: disable=protected-access
 
     except Exception as e:
         log.warning(f'Torch SDPA: {e}')
@@ -625,7 +625,23 @@ def set_cuda_params():
     except Exception:
         tunable = [False, False]
     log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} fp16={"pass" if fp16_ok else "fail"} bf16={"pass" if bf16_ok else "fail"} triton={"pass" if triton_ok else "fail"} optimization="{opts.cross_attention_optimization}"')
-    log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable}')
+    try:
+        num_threads = torch._inductor.config.compile_threads # pylint: disable=protected-access
+    except Exception:
+        num_threads = None
+    log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable} threads={num_threads}')
+
+    try:
+        from torch._inductor.runtime.runtime_utils import cache_dir
+        inductor_cache = cache_dir()
+    except Exception:
+        inductor_cache = os.getenv("TORCHINDUCTOR_CACHE_DIR", None)
+    try:
+        from triton import knobs
+        triton_cache = knobs.cache.dir
+    except Exception:
+        triton_cache = os.getenv("TRITON_CACHE_DIR", None)
+    log.info(f'Torch cache: inductor="{inductor_cache}" triton="{triton_cache}"')
 
 
 def randn(seed, shape=None):

modules/image/resize.py

@@ -44,6 +44,10 @@ def resize_image(resize_mode: int, im: Image.Image | torch.Tensor, width: int, h
             scale = max(w / im.width, h / im.height)
         if scale > 1.0:
             upscalers = [x for x in shared.sd_upscalers if x.name.lower().replace('-', ' ') == upscaler_name.lower().replace('-', ' ')]
+            if len(upscalers) == 0: # do force-refresh before failing
+                from modules.modelloader import load_upscalers
+                load_upscalers()
+                upscalers = [x for x in shared.sd_upscalers if x.name.lower().replace('-', ' ') == upscaler_name.lower().replace('-', ' ')]
             if len(upscalers) > 0:
                 selected_upscaler: upscaler.UpscalerData = upscalers[0]
                 if selected_upscaler.name.lower().startswith('latent'):
@@ -51,7 +55,7 @@ def resize_image(resize_mode: int, im: Image.Image | torch.Tensor, width: int, h
                 else:
                     im = selected_upscaler.scaler.upscale(im, scale, selected_upscaler.name)
             else:
-                log.warning(f"Resize upscaler: invalid={upscaler_name} fallback=resample")
+                log.warning(f'Resize upscaler: invalid="{upscaler_name}" fallback=resample')
                 log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}")
         if isinstance(im, Image.Image) and (im.width != w or im.height != h): # probably downsample after upscaler created larger image
             im = sharpfin.resize(im, (w, h))

modules/intel/ipex/__init__.py

@@ -55,7 +55,6 @@ def ipex_init(): # pylint: disable=too-many-statements
                 pass
             # Replace cuda with xpu:
             torch.cuda.current_device = torch.xpu.current_device
-            torch.cuda.current_stream = torch.xpu.current_stream
             torch.cuda.device = torch.xpu.device
             torch.cuda.device_count = torch.xpu.device_count
             torch.cuda.device_of = torch.xpu.device_of

modules/intel/ipex/hijacks.py

@@ -317,6 +317,14 @@ def torch_cuda_synchronize(device=None):
         return torch.xpu.synchronize(device)
 
 
+@wraps(torch.cuda.current_stream)
+def torch_cuda_current_stream(device=None):
+    if check_cuda(device):
+        return torch.xpu.current_stream(return_xpu(device))
+    else:
+        return torch.xpu.current_stream(device)
+
+
 @wraps(torch.cuda.device)
 def torch_cuda_device(device):
     if check_cuda(device):
@@ -342,7 +350,7 @@ def get_device_properties(device=None):
         "multi_processor_count": device_prop.gpu_subslice_count,
     }
     if not hasattr(device_prop, "L2_cache_size"):
-        new_keys["L2_cache_size"] = cache_size_dict.get(getattr(device_prop, "device_id", 0x56A0), cache_size_dict[0x0000])
+        new_keys["L2_cache_size"] = getattr(device_prop, "last_level_cache_size", cache_size_dict.get(getattr(device_prop, "device_id", 0x56A0), cache_size_dict[0x0000]))
     return DeviceProperties(device_prop, new_keys)
 
 
@@ -392,6 +400,7 @@ def ipex_hijacks():
     torch.load = torch_load
 
     torch.cuda.synchronize = torch_cuda_synchronize
+    torch.cuda.current_stream = torch_cuda_current_stream
     torch.cuda.device = torch_cuda_device
     torch.cuda.set_device = torch_cuda_set_device
     torch.cuda.get_device_properties = get_device_properties

modules/loader.py

@@ -56,6 +56,9 @@ except Exception as e:
     report(f'scipy=={scipy.__version__ if scipy is not None else None}', e)
 timer.startup.record("scipy")
 
+inductor_cache = os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "inductor"))
+triton_cache = os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "triton"))
+
 try:
     import atexit
     import torch._inductor.async_compile as ac

modules/logger.py

@@ -108,15 +108,35 @@ def setup_logging(debug=None, trace=None, filename=None):
             return ansi_escape.sub('', str(line))
 
         def emit(self, record):
-            if record.msg is not None and not isinstance(record.msg, str):
-                record.msg = str(record.msg)
+            if record.msg is None:
+                record.msg = ""
             try:
-                record.msg = record.msg.replace('"', "'")
+                msg = record.getMessage()
+            except Exception:
+                return
+            msg = msg.replace('"', "'")
+            msg = self.strip(msg)
+            try:
+                if '❱ ' in msg: # only last 3 lines of traceback
+                    lines = [l.strip() for l in msg.splitlines() if l.strip() and not l.startswith(' ')]
+                    if len(lines) > 3:
+                        lines = lines[-3:]
+                    lines = [l.replace('│   ', '').strip() for l in lines]
+                    lines.insert(0, 'Exception traceback:')
+                    msg = '\n'.join(lines)
+            except Exception:
+                pass
+            try:
+                if len(msg) > 1024:
+                    msg = msg[:1024] + '...'
+            except Exception:
+                pass
+            record.msg = msg
+            try:
+                formatted = self.format(record)
+                self.buffer.append(formatted)
             except Exception:
                 pass
-            line = self.format(record)
-            line = self.strip(line)
-            self.buffer.append(line[:1024])
             if len(self.buffer) > self.capacity:
                 self.buffer.pop(0)
 
@@ -128,7 +148,10 @@ def setup_logging(debug=None, trace=None, filename=None):
             super().__init__()
 
         def filter(self, record):
-            return len(record.getMessage()) > 2
+            try:
+                return len(record.getMessage()) > 2
+            except Exception:
+                return False
 
     def override_padding(self, console, options): # pylint: disable=redefined-outer-name
         style = console.get_style(self.style)

modules/lora/lora_load.py

@@ -58,8 +58,6 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
 
     sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
     cached = lora_cache.get(name, None)
-    if l.debug:
-        log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}')
     if cached is not None:
         return cached
     native_module = _NATIVE_DISPATCH.get(shared.sd_model_type)
@@ -279,7 +277,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         if network_on_disk is not None:
             shorthash = getattr(network_on_disk, 'shorthash', '').lower()
             if l.debug:
-                log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
+                log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}" cached={name in lora_cache}')
             try:
                 lora_scale = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
                 lora_module = lora_modules[i] if lora_modules and len(lora_modules) > i else None

modules/lora/native_adapter.py

@@ -56,9 +56,9 @@ KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_", "lor
 
 
 # Sentinel ``prefix_used`` value emitted by :func:`parse_key` when a bare path
-# starting with a member of ``bare_diffusers_prefixes`` matches. Loader
-# ``resolve_targets`` callables dispatch on this string to pass the base path
-# through verbatim (no rename required, the path is already in diffusers form).
+# starting with a member of ``bare_diffusers_prefixes`` matches. A loader
+# ``resolve_targets`` may dispatch on this string to rewrite the base path;
+# when it declines, :func:`resolve_group_targets` binds the path verbatim.
 BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers"
 
 
@@ -69,10 +69,12 @@ BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers"
 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. ``lycoris_`` bases are already
-# underscored, so the loader's ``.replace(".", "_")`` is a no-op on them.
+# Prefixes whose parsed ``base`` is normally already a network-key tail
+# (``arch_prefix + base.replace(".", "_")`` matches the stamped module name).
+# :func:`resolve_group_targets` binds them verbatim unless the arch's
+# ``resolve_targets`` claims the group first (needed when the arch's module
+# tree diverges from the names these formats carry). ``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_")
 
 
@@ -438,12 +440,15 @@ read_state_dict = sd_models.read_state_dict
 def resolve_group_targets(resolve_targets, prefix_used, base):
     """Map a parsed ``(prefix_used, base)`` group to ``[(diffusers_path, chunk), ...]``.
 
-    Passthrough prefixes (:data:`PASSTHROUGH_PREFIXES_DEFAULT`) bind verbatim;
-    everything else defers to the arch's ``resolve_targets``. Centralizing the
-    passthrough keeps each arch's resolver to the prefixes it actually rewrites.
+    The arch's ``resolve_targets`` is consulted first for every group.
+    Passthrough prefixes (:data:`PASSTHROUGH_PREFIXES_DEFAULT`) fall back to
+    verbatim binding when the resolver declines (returns empty), so arches
+    whose module tree matches the stamped names need no rewrite branch, while
+    arches with a divergent tree (e.g. krea2, whose transformer keeps checkpoint
+    names that differ from the ecosystem's diffusers names) can rewrite them.
     """
     if prefix_used in PASSTHROUGH_PREFIXES_DEFAULT:
-        return [(base, None)]
+        return resolve_targets(prefix_used, base) or [(base, None)]
     return resolve_targets(prefix_used, base)
 
 

modules/masking.py

@@ -264,7 +264,7 @@ def run_segment(input_image: gr.Image, input_mask: np.ndarray):
 def run_rembg(input_image: Image.Image, input_mask: np.ndarray):
     try:
         from installer import install
-        for pkg in ["dctorch==0.1.2", "pymatting", "pooch", "rembg"]:
+        for pkg in ["dctorch==0.1.2", "pymatting", "pooch", "rembg", "numba"]:
             install(pkg, no_deps=True, ignore=False)
         import rembg
     except Exception as e:

modules/model_quant.py

@@ -112,7 +112,10 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
                 quantized_matmul_dtype = shared.opts.sdnq_quantize_matmul_mode_te
             else:
                 quantized_matmul_dtype = shared.opts.sdnq_quantize_matmul_mode
-        if quantized_matmul_dtype == "auto":
+
+        use_quantized_matmul = quantized_matmul_dtype != "disabled"
+        quantized_matmul_dtype_log = quantized_matmul_dtype
+        if quantized_matmul_dtype in {"enabled", "disabled"}:
             quantized_matmul_dtype = None
 
         if modules_to_not_convert is None:
@@ -154,7 +157,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
             use_hadamard=shared.opts.sdnq_use_hadamard,
             quant_conv=shared.opts.sdnq_quantize_conv_layers,
             quant_embedding=shared.opts.sdnq_quantize_embedding_layers,
-            use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+            use_quantized_matmul=use_quantized_matmul,
             use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
             use_dynamic_quantization=shared.opts.sdnq_use_dynamic_quantization,
             dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
@@ -164,11 +167,9 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
             modules_to_not_convert=modules_to_not_convert,
             modules_dtype_dict=modules_dtype_dict.copy(),
         )
-        if quantized_matmul_dtype is None:
-            quantized_matmul_dtype = "auto" # set for logging
         svd = f'{shared.opts.sdnq_use_svd} rank={shared.opts.sdnq_svd_rank} steps={shared.opts.sdnq_svd_steps}' if shared.opts.sdnq_use_svd else f'{shared.opts.sdnq_use_svd}'
         hadamard = f'{shared.opts.sdnq_use_hadamard} group={shared.opts.sdnq_hadamard_group_size}' if shared.opts.sdnq_use_hadamard else f'{shared.opts.sdnq_use_hadamard}'
-        log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} svd={svd} hadamard={hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization} group={shared.opts.sdnq_group_size} loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_dtype={quantized_matmul_dtype} matmul_quant={shared.opts.sdnq_use_quantized_matmul} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embed={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} device={quantization_device} return={return_device} gpu={shared.opts.sdnq_quantize_with_gpu} map={shared.opts.device_map}')
+        log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} svd={svd} hadamard={hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization} group={shared.opts.sdnq_group_size} loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_dtype={quantized_matmul_dtype_log} matmul_quant={use_quantized_matmul} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embed={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} device={quantization_device} return={return_device} gpu={shared.opts.sdnq_quantize_with_gpu} map={shared.opts.device_map}')
         if len(modules_to_not_convert) > 0 or modules_dtype_dict:
             log.debug(f'Quantization: module={module} type=sdnq skip_modules={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
         if kwargs is None:
@@ -371,9 +372,13 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
             quantized_matmul_dtype = shared.opts.sdnq_quantize_matmul_mode_te
         else:
             quantized_matmul_dtype = shared.opts.sdnq_quantize_matmul_mode
-    if quantized_matmul_dtype == "auto":
+
+    use_quantized_matmul = quantized_matmul_dtype not in {"no", "disabled"}
+    quantized_matmul_dtype_log = quantized_matmul_dtype
+    if quantized_matmul_dtype in {"enabled", "disabled"}:
         quantized_matmul_dtype = None
 
+
     quantization_device, return_device = get_sdnq_devices(mode="post")
 
     if modules_to_not_convert is None:
@@ -416,7 +421,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
         use_hadamard=shared.opts.sdnq_use_hadamard,
         quant_conv=shared.opts.sdnq_quantize_conv_layers,
         quant_embedding=shared.opts.sdnq_quantize_embedding_layers,
-        use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+        use_quantized_matmul=use_quantized_matmul,
         use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
         use_dynamic_quantization=shared.opts.sdnq_use_dynamic_quantization,
         dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
@@ -457,10 +462,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
         model = model.to(devices.cpu)
     if do_gc:
         devices.torch_gc(force=True, reason='sdnq')
-
-    if quantized_matmul_dtype is None:
-        quantized_matmul_dtype = "auto" # set for logging
-    log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul_dtype={quantized_matmul_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} svd={shared.opts.sdnq_use_svd} hadamard={shared.opts.sdnq_use_hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization}:group={shared.opts.sdnq_group_size}:hadamard_group={shared.opts.sdnq_hadamard_group_size}:rank={shared.opts.sdnq_svd_rank}:steps={shared.opts.sdnq_svd_steps}:loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embedding={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} gpu={shared.opts.sdnq_quantize_with_gpu} device={quantization_device} return={return_device} map={shared.opts.device_map} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_skip={modules_to_not_convert} modules_dtype={modules_dtype_dict}')
+    log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul_dtype={quantized_matmul_dtype_log} matmul={use_quantized_matmul} svd={shared.opts.sdnq_use_svd} hadamard={shared.opts.sdnq_use_hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization}:group={shared.opts.sdnq_group_size}:hadamard_group={shared.opts.sdnq_hadamard_group_size}:rank={shared.opts.sdnq_svd_rank}:steps={shared.opts.sdnq_svd_steps}:loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embedding={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} gpu={shared.opts.sdnq_quantize_with_gpu} device={quantization_device} return={return_device} map={shared.opts.device_map} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_skip={modules_to_not_convert} modules_dtype={modules_dtype_dict}')
     return model
 
 

modules/modeldata.py

@@ -137,6 +137,8 @@ def get_model_type(pipe):
         model_type = 'hunyuanimage'
     elif 'sdxs-1b' in name:
         model_type = 'sdxs'
+    elif 'SeFi' in name:
+        model_type = 'sefi'
     # video models
     elif "Kandinsky5" in name and '2V' in name:
         model_type = 'kandinsky5video'

modules/modelloader.py

@@ -428,7 +428,7 @@ def move_files(src_path: str, dest_path: str, ext_filter: str | None = None):
         pass
 
 
-def load_upscalers():
+def load_upscalers(quiet=False):
     # We can only do this 'magic' method to dynamically load upscalers if they are referenced, so we'll try to import any _model.py files before looking in __subclasses__
     t0 = time.time()
     modules_dir = os.path.join(paths.script_path, "modules", "postprocess")
@@ -465,5 +465,6 @@ def load_upscalers():
         log.error('Upscalers: no data')
     shared.sd_upscalers = upscalers
     t1 = time.time()
-    log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={upscaler_types}")
+    if not quiet:
+        log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={upscaler_types}")
     return [x.name for x in shared.sd_upscalers]

modules/postprocess/seedvr_model.py

@@ -1,13 +1,15 @@
 import time
+import os
 import random
 import numpy as np
 import torch
 from PIL import Image
-from modules import devices
-from modules.shared import opts, log
+from modules import devices, timer
+from modules.shared import opts
 from modules.upscaler import Upscaler, UpscalerData
 from modules.image import convert
 from modules.model_quant import do_post_load_quant
+from modules.logger import log, console
 
 
 MODELS_MAP = {
@@ -29,6 +31,34 @@ class UpscalerSeedVR(Upscaler):
         self.model = None
         self.model_loaded = None
         self.device = devices.device
+        self.step = 1
+        self.frames = 0
+        self.offload = True
+        self.pbar = None
+        self.task = None
+        self.fps = 24
+        self.timer = None
+
+    def set_vae_params(self, vae_memory: float, tile_size: int, tile_overlap: float, vae_tile_encode: bool = True, vae_tile_decode: bool = True):
+        if vae_memory >= 0.99:
+            vae_memory = None
+        self.model.config.vae.memory_limit = {'conv_max_mem': vae_memory, 'norm_max_mem': vae_memory}
+        self.model.vae.set_memory_limit(**self.model.config.vae.memory_limit)
+        self.model.vae.tile_sample_min_size = tile_size
+        self.model.vae.tile_latent_min_size = tile_size // 8
+        self.model.vae.tile_overlap_factor = tile_overlap
+        if vae_tile_encode:
+            self.model.vae.use_slicing_encode = False
+            self.model.vae.use_tiling_encode = True
+        else:
+            self.model.vae.use_slicing_encode = True
+            self.model.vae.use_tiling_encode = False
+        if vae_tile_decode:
+            self.model.vae.use_slicing_decode = False
+            self.model.vae.use_tiling_decode = True
+        else:
+            self.model.vae.use_slicing_decode = True
+            self.model.vae.use_tiling_decode = False
 
     def load_model(self, path: str):
         model_name = MODELS_MAP.get(path, None)
@@ -43,68 +73,70 @@ class UpscalerSeedVR(Upscaler):
                 device=devices.device,
                 dtype=devices.dtype,
             )
-
             self.model_loaded = model_name
             self.model.dit.device = devices.device
             self.model.dit.dtype = devices.dtype
             self.model.vae_encode = self.vae_encode
             self.model.vae_decode = self.vae_decode
-            # Patch generation_loop's generation_step() with our wrapper; stash the original once
-            # so reloads don't re-wrap the wrapper itself (infinite recursion).
-            if not hasattr(generation, "generation_step_original"):
+            if not hasattr(generation, "generation_step_original"): # Patch generation_loop's generation_step() with our wrapper; stash the original once so reloads don't re-wrap the wrapper itself (infinite recursion).
                 generation.generation_step_original = generation.generation_step
             generation.generation_step = self.model_step
             self.model._internal_dict = {
                 'dit': self.model.dit,
                 'vae': self.model.vae,
             }
-            t1 = time.time()
             self.model.dit.config = self.model.config.dit
-            self.model.vae.tile_sample_min_size = 1024
-            self.model.vae.tile_latent_min_size = 128
-
             self.model = do_post_load_quant(self.model, allow=True)
-
+            t1 = time.time()
             log.info(f'Upscaler loaded: name="{self.name}" model="{model_name}" time={t1 - t0:.2f}')
 
     def vae_encode(self, samples):
-        log.debug(f'Upscaler encode: samples={samples[0].shape if len(samples) > 0 else None} tile={self.model.vae.tile_sample_min_size} overlap={self.model.vae.tile_overlap_factor}')
         latents = []
         if len(samples) == 0:
             return latents
-        self.model.dit = self.model.dit.to(device="cpu")
-        self.model.vae = self.model.vae.to(device=self.device)
-        devices.torch_gc()
+        self.pbar.update(self.task, description=f'encode: images={list(samples[0].shape) if len(samples) > 0 else None}')
+        if self.offload:
+            t0 = time.time()
+            self.model.dit = self.model.dit.to(device="cpu")
+            self.model.vae = self.model.vae.to(device=self.device)
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
+        t0 = time.time()
         from einops import rearrange
-        from modules.seedvr.src.optimization import memory_manager
-        memory_manager.clear_rope_cache(self.model)
         scale = self.model.config.vae.scaling_factor
         shift = self.model.config.vae.get("shifting_factor", 0.0)
         batches = [sample.unsqueeze(0) for sample in samples]
-        for sample in batches:
-            sample = sample.to(self.device, self.model.vae.dtype)
-            sample = self.model.vae.preprocess(sample)
-            latent = self.model.vae.encode(sample).latent
-            latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
-            latent = rearrange(latent, "b c ... -> b ... c")
-            latent = (latent - shift) * scale
-            latents.append(latent)
+        with devices.inference_context():
+            for sample in batches:
+                sample = sample.to(self.device, self.model.vae.dtype)
+                sample = self.model.vae.preprocess(sample)
+                latent = self.model.vae.encode(sample).latent
+                latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
+                latent = rearrange(latent, "b c ... -> b ... c")
+                latent = (latent - shift) * scale
+                latents.append(latent.contiguous())
         latents = [latent.squeeze(0) for latent in latents]
-        self.model.vae = self.model.vae.to(device="cpu")
-        devices.torch_gc()
+        self.timer.ts('encode', t0)
+        if self.offload:
+            t0 = time.time()
+            self.model.vae = self.model.vae.to(device="cpu")
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
         return latents
 
     def vae_decode(self, latents, target_dtype: torch.dtype = None):
-        log.debug(f'Upscaler decode: latents={latents[0].shape if len(latents) > 0 else None} tile={self.model.vae.tile_latent_min_size} overlap={self.model.vae.tile_overlap_factor}')
+        self.pbar.update(self.task, description=f'decode: latents={list(latents[0].shape) if len(latents) > 0 else None}')
         samples = []
         if len(latents) == 0:
             return samples
         from einops import rearrange
-        from modules.seedvr.src.optimization import memory_manager
-        memory_manager.clear_rope_cache(self.model)
-        self.model.dit = self.model.dit.to(device="cpu")
-        self.model.vae = self.model.vae.to(device=self.device)
-        devices.torch_gc()
+        if self.offload:
+            t0 = time.time()
+            self.model.dit = self.model.dit.to(device="cpu")
+            self.model.vae = self.model.vae.to(device=self.device)
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
+        t0 = time.time()
         scale = self.model.config.vae.scaling_factor
         shift = self.model.config.vae.get("shifting_factor", 0.0)
         latents = [latent.unsqueeze(0) for latent in latents]
@@ -113,59 +145,191 @@ class UpscalerSeedVR(Upscaler):
                 latent = latent.to(self.device, self.model.vae.dtype)
                 latent = latent / scale + shift
                 latent = rearrange(latent, "b ... c -> b c ...")
-                latent = latent.squeeze(2)
+                latent = latent.squeeze(2).contiguous()
                 sample = self.model.vae.decode(latent).sample
                 sample = self.model.vae.postprocess(sample)
-                samples.append(sample)
-        samples = [sample.squeeze(0) for sample in samples]
-        self.model.vae = self.model.vae.to(device="cpu")
-        devices.torch_gc()
+                samples.append(sample.squeeze(0).contiguous())
+        self.timer.ts('decode', t0)
+        if self.offload:
+            t0 = time.time()
+            self.model.vae = self.model.vae.to(device="cpu")
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
         return samples
 
     def model_step(self, *args, **kwargs):
+        from modules.shared import state
+        if state.interrupted or state.skipped:
+            return None
         from modules.seedvr.src.core import generation
-        from modules.seedvr.src.optimization import memory_manager
-        self.model.vae = self.model.vae.to(device="cpu")
-        self.model.dit = self.model.dit.to(device=self.device)
-        devices.torch_gc()
-        log.debug(f'Upscaler inference: args={len(args)} kwargs={list(kwargs.keys())}')
-        memory_manager.preinitialize_rope_cache(self.model)
-        with devices.inference_context():
-            result = generation.generation_step_original(*args, **kwargs)
-        self.model.dit = self.model.dit.to(device="cpu")
-        devices.torch_gc()
-        return result
-
-    def do_upscale(self, img: Image.Image, selected_file):
-        self.load_model(selected_file)
-        if self.model is None:
-            return img
-
-        from modules.seedvr.src.core import generation
-
-        width = int(self.scale * img.width) // 8 * 8
-        image_tensor = np.array(img)
-        image_tensor = torch.from_numpy(image_tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0
-
-        random.seed()
-        seed = int(random.randrange(4294967294))
-
+        if self.offload:
+            t0 = time.time()
+            self.model.vae = self.model.vae.to(device="cpu")
+            self.model.dit = self.model.dit.to(device=self.device)
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
         t0 = time.time()
         with devices.inference_context():
+            self.pbar.update(self.task, description=f'inference: batch={self.step}')
+            result = generation.generation_step_original(*args, **kwargs)
+            self.pbar.update(self.task, advance=self.step)
+        self.timer.ts('step', t0)
+        if self.offload:
+            t0 = time.time()
+            self.model.dit = self.model.dit.to(device="cpu")
+            self.timer.ts('offload', t0)
+        devices.torch_gc(fast=True)
+        return result
+
+    def read_image(self, image: str | Image.Image):
+        try:
+            if isinstance(image, str):
+                image = Image.open(image)
+            image = image.convert("RGB")
+            width = image.width
+            tensor = np.array(image)
+            tensor = torch.from_numpy(tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0
+            self.frames = 1
+            return tensor, width
+        except Exception as e:
+            log.error(f'Upscaler: name="SeedVR2" image="{image}" {e}')
+            return None, None
+
+    def read_video(self, video_path: str):
+        try:
+            import cv2
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                log.error(f'Upscaler: name="SeedVR2" video="{video_path}" failed to open')
+                return None, None
+            frames = []
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            self.fps = int(cap.get(cv2.CAP_PROP_FPS))
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                frames.append(frame)
+            cap.release()
+            if len(frames) == 0:
+                log.error(f'Upscaler: name="SeedVR2" video="{video_path}" no frames read')
+                return None, None
+            tensor = torch.from_numpy(np.array(frames)).to(device=devices.device, dtype=devices.dtype) / 255.0
+            self.frames = tensor.shape[0]
+            return tensor, width
+        except Exception as e:
+            log.error(f'Upscaler: name="SeedVR2" video="{video_path}" {e}')
+            return None, None
+
+    def create_video(self, tensor: torch.Tensor, codec: str = 'libx264', codec_opt: str = 'crf:16', interpolate: int = 0):
+        t0 = time.time()
+        from modules.video_models.video_save import save_video
+        pixels = tensor.permute(3, 0, 1, 2).unsqueeze(0) # from (t, h, w, c) to (n, c, t, h, w)
+        _frames, filename, _thumb = save_video(p=None,
+                                               pixels=pixels,
+                                               mp4_fps=self.fps,
+                                               mp4_thumb=False,
+                                               mp4_frames=False,
+                                               reclamp=False,
+                                               mp4_codec=codec,
+                                               mp4_opt=codec_opt,
+                                               mp4_interpolate=interpolate,
+                                              )
+        self.timer.ts('save', t0)
+        return filename
+
+    def do_upscale(self,
+                   img: Image.Image | str,
+                   selected_file,
+                   cfg_scale: float = 1.5,
+                   cfg_rescale: float = 0.0,
+                   steps: int = 1,
+                   seed: int = -1,
+                   scale: float | None = None,
+                   tile_size: int = 1024,
+                   tile_overlap: float = 0.25,
+                   batch_size: int = 1,
+                   batch_overlap: int = 0,
+                   offload: bool = True,
+                   interpolate: int = 1,
+                   codec: str = 'libx264',
+                   codec_opt: str = 'crf:16',
+                   vae_memory: float = 0.2,
+                   vae_tile_encode: bool = True,
+                   vae_tile_decode: bool = True,
+                  ):
+        self.timer = timer.Timer()
+        self.offload = offload
+        self.load_model(selected_file)
+        self.set_vae_params(vae_memory=vae_memory, tile_size=tile_size, tile_overlap=tile_overlap, vae_tile_encode=vae_tile_encode, vae_tile_decode=vae_tile_decode)
+        if self.model is None:
+            return img
+        if not self.offload:
+            self.model.dit = self.model.dit.to(device=devices.device)
+            self.model.vae = self.model.vae.to(device=devices.device)
+            devices.torch_gc(fast=True)
+        self.timer.record('load')
+
+        from modules.seedvr.src.core import generation
+
+        self.scale = self.scale if scale is None else scale
+        if isinstance(img, Image.Image):
+            tensor, width = self.read_image(img)
+        elif isinstance(img, str):
+            tensor, width = self.read_video(img)
+        else:
+            log.error(f'Upscaler: name="SeedVR2" image="{img}" unsupported type {type(img)}')
+            return img
+        self.timer.record('read')
+
+        if tensor is None or width is None:
+            log.error(f'Upscaler: name="SeedVR2" image="{img}" failed to read')
+            return img
+        width = int(self.scale * width) // 8 * 8
+        random.seed()
+        seed = int(random.randrange(4294967294)) if seed == -1 else int(seed)
+        self.step = 1 if self.frames == 1 else batch_size - batch_overlap
+
+        mode = "mode=image" if self.frames == 1 else f"mode=video frames={self.frames}"
+        batch_info = f'batch=(size={batch_size} overlap={batch_overlap})'
+        vae_info = f'vae=(tiled={vae_tile_encode}/{vae_tile_decode} memory={vae_memory} size={tile_size} overlap={tile_overlap})'
+        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" {mode} scale={self.scale} cfg={cfg_scale}:{cfg_rescale} seed={seed} steps={steps} offload={self.offload} {batch_info} {vae_info}')
+
+        import rich.progress as rp
+        self.pbar = rp.Progress(rp.TextColumn('[cyan]SeedVR:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console)
+        self.task = self.pbar.add_task(total=self.frames, description='starting...')
+        with devices.inference_context(), self.pbar:
+            self.pbar.update(self.task, description='initialize rope')
+            from modules.seedvr.src.optimization import memory_manager
+            memory_manager.clear_rope_cache(self.model)
+            memory_manager.preinitialize_rope_cache(self.model)
+            self.timer.record('init')
             result_tensor = generation.generation_loop(
                 runner=self.model,
-                images=image_tensor,
-                cfg_scale=opts.seedvr_cfg_scale,
+                images=tensor,
+                cfg_scale=cfg_scale,
+                cfg_rescale=cfg_rescale,
+                steps=steps, # TODO SeedVR steps
+                batch_size=batch_size, # TODO SeedVR batch size
+                temporal_overlap=batch_overlap, # TODO SeedVR temporal overlap
                 seed=seed,
                 res_w=width,
-                batch_size=1,
-                temporal_overlap=0,
                 device=devices.device,
+                color_reconstruct=True,
             )
-        t1 = time.time()
-        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={opts.seedvr_cfg_scale} seed={seed} time={t1 - t0:.2f}')
-        img = convert.to_pil(result_tensor.squeeze())
+            memory_manager.clear_rope_cache(self.model)
 
+        self.pbar.update(self.task, completed=self.frames)
+        t1 = time.time()
+        self.frames = result_tensor.shape[0] if result_tensor is not None else 0
+        self.timer.add('inference', self.timer.get('step') - self.timer.get('encode') - self.timer.get('decode'))
+        self.timer.rm('step')
+
+        t0 = time.time()
+        self.model.dit = self.model.dit.to(device="cpu")
+        self.model.vae = self.model.vae.to(device="cpu")
+        self.timer.ts('offload', t0)
         if opts.upscaler_unload:
             self.model.dit = None
             self.model.vae = None
@@ -173,4 +337,14 @@ class UpscalerSeedVR(Upscaler):
             self.model = None
             log.debug(f'Upscaler unload: type="{self.name}" model="{selected_file}"')
         devices.torch_gc(force=True)
-        return img
+        self.timer.ts('cleanup', t1)
+
+        if self.frames == 1:
+            result = convert.to_pil(result_tensor.squeeze())
+        elif self.frames > 1:
+            result = self.create_video(result_tensor, codec=codec, codec_opt=codec_opt, interpolate=interpolate)
+        else:
+            log.error(f'Upscaler: name="SeedVR2" model="{selected_file}" no frames generated')
+            result = img
+        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" frames={self.frames} {self.timer.summary()}')
+        return result

modules/postprocessing.py

@@ -3,104 +3,164 @@ import tempfile
 
 from PIL import Image
 
-from modules import shared, images, devices, scripts_manager, scripts_postprocessing, infotext
+from modules import shared, images, devices, errors, scripts_manager, scripts_postprocessing, infotext
 from modules.logger import log
 from modules.shared import opts
 from modules.paths import resolve_output_path
 
 
-def run_postprocessing(extras_mode, image, image_folder: list[tempfile.NamedTemporaryFile], input_dir, output_dir, show_extras_results, *args, save_output: bool = True):
+def run_postprocessing(extras_mode,
+                       image,
+                       image_folder: list[tempfile.NamedTemporaryFile],
+                       input_dir,
+                       output_dir,
+                       video,
+                       show_extras_results,
+                       *args,
+                       save_output: bool = True):
     devices.torch_gc()
-    shared.state.begin('Extras')
-    image_data = []
-    image_names = []
-    image_fullnames = []
-    image_ext = []
-    outputs = []
-    params = {}
-    info = ''
-    if extras_mode == 1:
-        for img in image_folder:
-            if isinstance(img, Image.Image):
-                image = img
-                fn = ''
-                ext = None
-            else:
+    job_id = shared.state.begin('Process')
+
+    def prepare_inputs(image):
+        image_data = []
+        image_names = []
+        image_ext = []
+        if extras_mode == 1: # process batch
+            for img in image_folder:
+                if isinstance(img, Image.Image):
+                    image = img
+                    fn = ''
+                    ext = None
+                else:
+                    try:
+                        image = Image.open(os.path.abspath(img.name))
+                    except Exception as e:
+                        log.error(f'Failed to open image: file="{img.name}" {e}')
+                        continue
+                    fn, ext = os.path.splitext(img.orig_name)
+                image_data.append(image)
+                image_names.append(fn)
+                image_ext.append(ext)
+            log.debug(f'Process: mode=batch inputs={len(image_folder)} images={len(image_data)}')
+        elif extras_mode == 2: # process folder
+            assert input_dir, 'input directory not selected'
+            image_list = os.listdir(input_dir)
+            for filename in image_list:
+                fn = os.path.join(input_dir, filename)
                 try:
-                    image = Image.open(os.path.abspath(img.name))
+                    image = Image.open(fn)
                 except Exception as e:
-                    log.error(f'Failed to open image: file="{img.name}" {e}')
+                    log.error(f'Failed to open image: file="{fn}" {e}')
                     continue
-                fn, ext = os.path.splitext(img.orig_name)
-                image_fullnames.append(img.name)
+                image_data.append(image)
+                image_names.append(fn)
+                image_ext.append(None)
+            log.debug(f'Process: mode=folder inputs={input_dir} files={len(image_list)} images={len(image_data)}')
+        elif extras_mode == 3: # process video
+            pass
+        else: # process image
             image_data.append(image)
-            image_names.append(fn)
-            image_ext.append(ext)
-        log.debug(f'Process: mode=batch inputs={len(image_folder)} images={len(image_data)}')
-    elif extras_mode == 2:
-        assert input_dir, 'input directory not selected'
-        image_list = os.listdir(input_dir)
-        for filename in image_list:
-            fn = os.path.join(input_dir, filename)
-            try:
-                image = Image.open(fn)
-            except Exception as e:
-                log.error(f'Failed to open image: file="{fn}" {e}')
-                continue
-            image_fullnames.append(fn)
-            image_data.append(image)
-            image_names.append(fn)
+            image_names.append(None)
             image_ext.append(None)
-        log.debug(f'Process: mode=folder inputs={input_dir} files={len(image_list)} images={len(image_data)}')
-    else:
-        image_data.append(image)
-        image_names.append(None)
-        image_ext.append(None)
+        return image_data, image_names, image_ext
+
+    image_data, image_names, image_ext = prepare_inputs(image)
+
     if extras_mode == 2 and output_dir != '':
         outpath = output_dir
     else:
         outpath = resolve_output_path(opts.outdir_samples, opts.outdir_extras_samples)
-    processed_images = []
-    for image, name, ext in zip(image_data, image_names, image_ext, strict=False): # pylint: disable=redefined-argument-from-local
-        log.debug(f'Process: image={image} {args}')
+
+    def process_images():
+        outputs = []
+        params = {}
         info = ''
-        if shared.state.interrupted:
-            log.debug('Postprocess interrupted')
-            break
-        if image is None:
-            continue
-        shared.state.textinfo = name
-        pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))
+        processed_images = []
+        for image, name, ext in zip(image_data, image_names, image_ext, strict=False): # pylint: disable=redefined-argument-from-local
+            log.debug(f'Process: image={image} {args}')
+            info = ''
+            if shared.state.interrupted:
+                log.debug('Postprocess interrupted')
+                break
+            if isinstance(image, str):
+                try:
+                    image = Image.open(image)
+                except Exception as e:
+                    log.error(f'Failed to open image: file="{image}" {e}')
+                    continue
+            if image is None:
+                continue
+            shared.state.textinfo = name
+            pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))
+            scripts_manager.scripts_postproc.run(pp, args)
+            geninfo, items = images.read_info_from_image(image)
+            params = infotext.parse(geninfo)
+            for k, v in items.items():
+                pp.image.info[k] = v
+            if 'parameters' in items:
+                info = items['parameters'] + ', '
+            if (params.get('size-1', 0) != pp.image.width) or (params.get('size-2', 0) != pp.image.height):
+                params['size-1'] = pp.image.width
+                params['size-2'] = pp.image.height
+                info += f"Size: {pp.image.width}x{pp.image.height}, "
+            info = info + ", ".join([k if k == v else f'{k}: {infotext.quote(v)}' for k, v in pp.info.items() if v is not None])
+            pp.image.info["postprocessing"] = info
+            processed_images.append(pp.image)
+            if save_output:
+                if opts.use_original_name_batch and name is not None:
+                    forced_filename = os.path.splitext(os.path.basename(name))[0]
+                    images.save_image(pp.image, path=outpath, extension=ext or opts.samples_format, info=info, grid=False, pnginfo_section_name="extras", existing_info=pp.image.info, forced_filename=forced_filename)
+                else:
+                    images.save_image(pp.image, path=outpath, extension=ext or opts.samples_format, info=info, grid=False, pnginfo_section_name="extras", existing_info=pp.image.info)
+            if extras_mode != 2 or show_extras_results:
+                outputs.append(pp.image)
+            image.close()
+        scripts_manager.scripts_postproc.postprocess(processed_images, args)
+        return outputs, info, params
+
+    def process_video():
+        outputs = []
+        params = {}
+        if not video or not isinstance(video, str) or not os.path.isfile(video):
+            log.error(f'Process: mode=video file="{video}" not found')
+            return outputs, video, '', params
+        log.debug(f'Process: video={video} {args}')
+        shared.state.textinfo = video
+        pp = scripts_postprocessing.PostprocessedImage(video=video)
         scripts_manager.scripts_postproc.run(pp, args)
-        geninfo, items = images.read_info_from_image(image)
-        params = infotext.parse(geninfo)
-        for k, v in items.items():
-            pp.image.info[k] = v
-        if 'parameters' in items:
-            info = items['parameters'] + ', '
-        if (params.get('size-1', 0) != pp.image.width) or (params.get('size-2', 0) != pp.image.height):
-            params['size-1'] = pp.image.width
-            params['size-2'] = pp.image.height
-            info += f"Size: {pp.image.width}x{pp.image.height}, "
-        info = info + ", ".join([k if k == v else f'{k}: {infotext.quote(v)}' for k, v in pp.info.items() if v is not None])
-        pp.image.info["postprocessing"] = info
-        processed_images.append(pp.image)
-        if save_output:
-            if opts.use_original_name_batch and name is not None:
-                forced_filename = os.path.splitext(os.path.basename(name))[0]
-                images.save_image(pp.image, path=outpath, extension=ext or opts.samples_format, info=info, grid=False, pnginfo_section_name="extras", existing_info=pp.image.info, forced_filename=forced_filename)
-            else:
-                images.save_image(pp.image, path=outpath, extension=ext or opts.samples_format, info=info, grid=False, pnginfo_section_name="extras", existing_info=pp.image.info)
-        if extras_mode != 2 or show_extras_results:
-            outputs.append(pp.image)
-        image.close()
-    scripts_manager.scripts_postproc.postprocess(processed_images, args)
+
+        from modules.video import get_video_info
+        params = get_video_info(pp.video)
+        info = ', '.join([f'{k}: {v}' for k, v in params.items()])
+        return pp.video, info, params
+
+    if extras_mode == 3:
+        try:
+            video, info, params = process_video()
+        except Exception as e:
+            log.error(f'Process: mode=video {e}')
+            errors.display(e, 'postprocessing video')
+            video = None
+            info = str(e)
+            params = {}
+        outputs = []
+    else:
+        try:
+            outputs, info, params = process_images()
+        except Exception as e:
+            log.error(f'Process: mode=image {e}')
+            errors.display(e, 'postprocessing image')
+            outputs = []
+            info = str(e)
+            params = {}
+        video = None
 
     devices.torch_gc()
-    return outputs, info, params
+    shared.state.end(job_id)
+    return outputs, video, info, params
 
 
-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, save_output: bool = True, script_args: dict | None = None):
+def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, video, show_extras_results, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, save_output: bool = True, script_args: dict | None = None):
     """old handler for API"""
 
     merged = {
@@ -120,4 +180,4 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
             merged.setdefault(name, {}).update(kvs or {})
     args = scripts_manager.scripts_postproc.create_args_for_run(merged)
 
-    return run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output=save_output)
+    return run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, video, show_extras_results, *args, save_output=save_output)

modules/processing.py

@@ -416,6 +416,23 @@ def process_samples(p: StableDiffusionProcessing, samples):
     return out_images, out_infotexts
 
 
+def print_stats():
+    log.debug(f'Processed: timers={timer.process.dct()}')
+    log.debug(f'Processed: memory={memstats.memory_stats()}')
+
+    if shared.opts.sdnq_dequantize_compile:
+        from modules.timer_sdnq import update_sdnq_attention_timers
+        update_sdnq_attention_timers()
+        if timer.autotune.get_total() > 0.001:
+            log.debug(f'Processed: autotune={timer.autotune.dct(min_time=0)}')
+
+    if devices.triton_ok:
+        from modules.sd_models_compile import update_compile_times
+        update_compile_times()
+        if timer.dynamo.get_total() > 0.001:
+            log.debug(f'Processed: dynamo={timer.dynamo.dct(min_time=2.0, no_total=True)}')
+
+
 def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if type(p.prompt) == list:
         assert len(p.prompt) > 0
@@ -570,10 +587,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         p.scripts.postprocess(p, results)
     timer.process.record('post')
     p.ops = list(set(p.ops))
+
     if not p.disable_extra_networks:
         log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} ops={p.ops}')
-        log.debug(f'Processed: timers={timer.process.dct()}')
-        log.debug(f'Processed: memory={memstats.memory_stats()}')
+        print_stats()
 
     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
         devices.torch_gc(force=True, reason='final')

modules/processing_diffusers.py

@@ -542,7 +542,8 @@ def update_pipeline(sd_model, p: processing.StableDiffusionProcessing):
         updated_model = preprocess_onnx_pipeline(p)
         global orig_pipeline # pylint: disable=global-statement
         orig_pipeline = updated_model # processed ONNX pipeline should not be replaced with original pipeline.
-    if getattr(updated_model, "current_attn_name", None) != shared.opts.cross_attention_optimization:
+    current_attn = getattr(updated_model, "current_attn_name", None)
+    if (current_attn != shared.opts.cross_attention_optimization) and (current_attn != shared.opts.sdp_overrides):
         log.info(f"Setting attention optimization: {shared.opts.cross_attention_optimization}")
         attention.set_diffusers_attention(updated_model)
     return updated_model

modules/processing_prompt.py

@@ -23,7 +23,7 @@ def fix_prompt_batch(p, prompts, negative_prompts, prompts_2, negative_prompts_2
     if type(negative_prompts) is str:
         negative_prompts = [negative_prompts]
 
-    if hasattr(p, 'init_images') and p.init_images is not None and len(p.init_images) > 1:
+    if hasattr(p, 'init_images') and (p.init_images is not None) and (len(p.init_images) > 1) and not getattr(p, 'skip_processing', False):
         while len(prompts) < len(p.init_images):
             prompts.append(prompts[-1] if prompts else '')
         while len(negative_prompts) < len(p.init_images):

modules/rembg/rembg_api.py

@@ -1,6 +1,12 @@
 from fastapi import Body
 
 
+def dependencies():
+    from installer import install
+    for pkg in ["dctorch==0.1.2", "pymatting", "pooch", "rembg", "numba"]:
+        install(pkg, no_deps=True, ignore=False)
+
+
 async def post_rembg(
     input_image: str = Body("", title='rembg input image'),
     model: str = Body("u2net", title='rembg model'),
@@ -23,9 +29,7 @@ async def post_rembg(
         from modules.rembg import ben2
         image = ben2.remove(input_image, refine=refine)
     else:
-        from installer import install
-        for pkg in ["dctorch==0.1.2", "pymatting", "pooch", "rembg"]:
-            install(pkg, no_deps=True, ignore=False)
+        dependencies()
         import rembg
         image = rembg.remove( # pylint: disable=c-extension-no-member
             input_image,

modules/scripts_postprocessing.py

@@ -5,14 +5,15 @@ from modules.logger import log
 
 
 class PostprocessedImage:
-    def __init__(self, image, info = None):
+    def __init__(self, image = None, video = None, info = None):
         if info is None:
             info = {}
         self.image = image
+        self.video = video
         self.info = info
 
     def __str__(self):
-        return f'PostprocessedImage(image={self.image} info={self.info})'
+        return f'PostprocessedImage(image={self.image} video={self.video} info={self.info})'
 
 
 class ScriptPostprocessing:
@@ -160,6 +161,6 @@ class ScriptPostprocessingRunner:
             else:
                 for (name, _component), value in zip(script.controls.items(), script_args, strict=False):
                     process_kwargs[name] = value
-            log.debug(f'Postprocess: script={script.name} args={process_args} kwargs={process_kwargs}')
+            log.debug(f'Postprocess: script="{script.name}" args={process_args} kwargs={process_kwargs}')
             script.postprocess(filenames, *process_args, **process_kwargs)
             shared.state.end(jobid)

modules/sd_detect.py

@@ -173,6 +173,8 @@ def guess_by_name(fn, current_guess):
         new_guess = 'VIBE'
     elif 'joyai-image-edit' in fn.lower() or 'joy-image-edit' in fn.lower():
         new_guess = 'JoyEdit'
+    elif 'sefi-image' in fn.lower():
+        new_guess = 'SeFi'
     if debug_load:
         log.trace(f'Autodetect: method=name file="{fn}" previous="{current_guess}" current="{new_guess}"')
     return new_guess or current_guess

modules/sd_hijack_te.py

@@ -80,6 +80,8 @@ def hijack_encode_prompt(*args, **kwargs):
         errors.display(e, 'Encode prompt')
     t1 = time.time()
     timer.process.add('te', t1-t0)
+    if t1 - t0 > 10:
+        log.warning(f'Encode: time={t1-t0:.3f} long encode prompt')
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     shared.state.end(jobid)
     # from modules import memstats

modules/sd_models.py

@@ -194,10 +194,17 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
     for module_name in get_module_names(sd_model):
         module = getattr(sd_model, module_name, None)
         if hasattr(module, "quantization_config") and getattr(module.quantization_config, "quant_method", None) == "sdnq":
-            if module.quantization_config.use_quantized_matmul != shared.opts.sdnq_use_quantized_matmul:
+            if module_name.startswith("text_encoder"):
+                if shared.opts.sdnq_quantize_matmul_mode_te == "Same as model":
+                    sdnq_use_quantized_matmul = shared.opts.sdnq_quantize_matmul_mode != "disabled"
+                else:
+                    sdnq_use_quantized_matmul = shared.opts.sdnq_quantize_matmul_mode_te != "disabled"
+            else:
+                sdnq_use_quantized_matmul = shared.opts.sdnq_quantize_matmul_mode != "disabled"
+            if module.quantization_config.use_quantized_matmul != sdnq_use_quantized_matmul:
                 from modules.sdnq.loader import apply_sdnq_options_to_model
-                # log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={shared.opts.sdnq_use_quantized_matmul}')
-                module = apply_sdnq_options_to_model(module, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul)
+                # log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={sdnq_use_quantized_matmul}')
+                module = apply_sdnq_options_to_model(module, use_quantized_matmul=sdnq_use_quantized_matmul)
                 setattr(sd_model, module_name, module)
 
     if offload:
@@ -595,6 +602,10 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf
             from pipelines.model_sdxs import load_sdxs
             sd_model = load_sdxs(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['SeFi']:
+            from pipelines.model_sefi import load_sefi
+            sd_model = load_sefi(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
     except Exception as e:
         log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
         errors.display(e, 'Load')

modules/sd_models_compile.py

@@ -370,3 +370,38 @@ def openvino_post_compile(op="base"): # delete unet after OpenVINO compile
             if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_refiner, "unet"):
                 shared.sd_refiner.unet.apply(sd_models_utils.convert_to_faketensors)
                 devices.torch_gc(force=True)
+
+
+def update_compile_times():
+    from modules.timer import dynamo
+    dynamo.reset()
+    try:
+        from torch._dynamo.utils import compile_times, reset_frame_count
+        raw_str = compile_times()
+        reset_frame_count()
+    except Exception:
+        return
+    lines = raw_str.strip().split('\n')
+    # parsed = []
+    for line in lines:
+        if not line or 'TorchDynamo compilation metrics' in line or 'Function, Runtimes' in line:
+            continue
+        parts = line.split(',')
+        fn = parts[0].strip()
+        try:
+            times = [float(t.strip()) for t in parts[1:] if t.strip()]
+            if times:
+                # parsed.append((fn, sum(times), len(times), max(times)))
+                dynamo.add(fn, round(sum(times), 2))
+        except ValueError:
+            continue
+    """
+    parsed.sort(key=lambda x: x[1], reverse=True)
+    results = {}
+    min_time = 0.1
+    for fn, total, count, max_val in parsed:
+        if total > min_time:
+            dynamo.ts(fn, total)
+            results[fn] = { "total": round(total, 2), "count": count, "avg": round(total / count, 2), "max": round(max_val, 2) }
+    return results
+    """

modules/sdnq/common.py

@@ -6,7 +6,7 @@ import torch
 
 from modules import shared
 
-sdnq_version = "0.2.2"
+sdnq_version = "0.2.3"
 sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"}
 
 torch_version = torch.__version__[:4]
@@ -338,9 +338,11 @@ weights_dtype_order = [
 
 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 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

@@ -13,10 +13,6 @@ 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,
@@ -300,7 +296,12 @@ class SDNQDequantizer:
     ) -> 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)
-        re_quantize_matmul_func = re_quantize_matmul if skip_compile or skip_fp8_compile(self.weights_dtype) else re_quantize_matmul_compiled
+        if skip_compile:
+            re_quantize_matmul_func = re_quantize_matmul
+        else:
+            re_quantize_matmul_func = re_quantize_matmul_compiled
+            if not is_fp8_compile_supported and weight.dtype == torch.float8_e4m3fn:
+                weight = weight.to(dtype=scale.dtype)
         return re_quantize_matmul_func(
             self.weights_dtype,
             weight,
@@ -333,7 +334,12 @@ 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 or skip_fp8_compile(self.weights_dtype) else dequantize_weight_compiled
+        if skip_compile:
+            dequantize_weight_func = dequantize_weight
+        else:
+            dequantize_weight_func = dequantize_weight_compiled
+            if not is_fp8_compile_supported and weight.dtype == torch.float8_e4m3fn:
+                weight = weight.to(dtype=scale.dtype)
         return dequantize_weight_func(
             self.weights_dtype,
             weight,

modules/sdnq/kernel_wrappers.py

@@ -4,7 +4,8 @@ import os
 import sys
 import torch
 
-from modules import devices
+from modules import devices, shared
+from .common import compile_func
 
 
 if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
@@ -32,15 +33,21 @@ if devices.backend == "rocm":
 else:
     is_rdna2_and_older = False
 
+if devices.backend in {"ipex", "xpu"}:
+    is_alchemist_or_igpu = bool(not torch.xpu.get_device_capability(devices.device).get("has_subgroup_2d_block_io", False))
+else:
+    is_alchemist_or_igpu = 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"})
+    use_triton_mm = bool(not is_alchemist_or_igpu and (devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"}))
 else:
     use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"})
+use_triton_scaled_mm = bool(use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").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
@@ -51,9 +58,11 @@ else:
 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")
+    use_contiguous_fp8_mm = use_contiguous_fp16_mm
 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
+    use_contiguous_fp8_mm = use_contiguous_fp16_mm
 
 
 int_mm_func = None
@@ -63,6 +72,7 @@ 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
@@ -73,26 +83,57 @@ if use_openvino_mm:
 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_tensorwise_fp8_matmul = True
+        if use_triton_scaled_mm:
+            from .kernels.triton_scaled_mm import sdnq_scaled_mm
+            int_scaled_mm_func = sdnq_scaled_mm
+            fp_scaled_mm_func = sdnq_scaled_mm
+            if is_fp8_mm_supported:
+                fp8_scaled_mm_func = sdnq_scaled_mm
+    except Exception as e:
         use_triton_mm = False
+        use_triton_scaled_mm = False
+        shared.log.warning(f"SDNQ: Triton kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
 
-if fp_mm_func is None and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
+
+if (
+    fp_mm_func is None and not is_alchemist_or_igpu
+    and devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"}
+    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
+        if use_triton_scaled_mm:
+            from .kernels.triton_scaled_mm import sdnq_scaled_mm
+            fp_scaled_mm_func = sdnq_scaled_mm
     except Exception:
         use_triton_mm = False
+        use_triton_scaled_mm = False
+
+
+if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None:
+    include_mm_kernel_in_compile = bool(not use_triton_scaled_mm)
+else:
+    include_mm_kernel_in_compile = bool(os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", "0").lower() not in {"0", "false", "no"})
+
+
+def fp_mm_torch_cuda(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+    return torch.mm(a,b, out_dtype=out_dtype)
+
+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)
 
 
 if int_mm_func is None:
@@ -100,47 +141,45 @@ if int_mm_func is None:
         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)
+        fp_mm_func = fp_mm_torch_cuda
     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
+        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 is_fp8_mm_supported:
+        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
+        use_contiguous_fp8_mm = False
+    else:
+        fp8_mm_func = fp_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)
+            return int_mm_func(a,b, out_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
+            return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+    int_scaled_mm_func = compile_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)
+            return fp_mm_func(a,b, out_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
+            return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+    fp_scaled_mm_func = compile_func(fp_scaled_mm_torch)
+
 
 if fp8_scaled_mm_func is None:
-    if use_tensorwise_fp8_matmul:
+    if use_tensorwise_fp8_matmul or not is_fp8_mm_supported:
         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)
@@ -152,4 +191,4 @@ if fp8_scaled_mm_func is None:
                 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
+    fp8_scaled_mm_func = compile_func(fp8_scaled_mm_torch)

modules/sdnq/kernels/triton_atten.py

@@ -14,8 +14,8 @@ matmul_configs = [
     triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN}, num_warps=w, num_stages=s)
     for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
     for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST", "32,64").replace(" ","").split(",")]
-    for w in [int(w) for w in os.environ.get("SDNQ_TRITON_ATTEN_NUM_WARPS_LIST", "4,8").replace(" ","").split(",")]
-    for s in [int(s) for s in os.environ.get("SDNQ_TRITON_ATTEN_NUM_STAGES_LIST", "1,2,4").replace(" ","").split(",")]
+    for w in [int(w) for w in os.environ.get("SDNQ_TRITON_ATTEN_NUM_WARPS_LIST", "8,16" if torch.xpu.is_available() else "4,8").replace(" ","").split(",")]
+    for s in [int(s) for s in os.environ.get("SDNQ_TRITON_ATTEN_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
 ]
 
 
@@ -91,7 +91,7 @@ def sdnq_attn_kernel(
     tl.assume(qk_is_quantized == 0 or qk_is_quantized == 1) # pylint: disable=consider-using-in
     tl.assume(pv_is_quantized == 0 or pv_is_quantized == 1) # pylint: disable=consider-using-in
 
-    do_k_mask: tl.constexpr = KN % BLOCK_SIZE_N != 0
+    do_k_mask = KN % BLOCK_SIZE_N != 0
     start_m_block = start_m * BLOCK_SIZE_M
     offs_m = start_m_block + tl.arange(0, BLOCK_SIZE_M)
     offs_n = tl.arange(0, BLOCK_SIZE_N)
@@ -139,9 +139,9 @@ def sdnq_attn_kernel(
             if qk_is_quantized:
                 k_scale = k_scale_desc.load([start_n])[None, :]
                 if q.dtype == tl.int8:
-                    qk = tl.dot(q, k, out_dtype=tl.int32).to(tl.float32) * q_scale * k_scale
+                    qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.int32).to(tl.float32), q_scale), k_scale)
                 else:
-                    qk = tl.dot(q, k, out_dtype=tl.float32) * q_scale * k_scale
+                    qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.float32), q_scale), k_scale)
             else:
                 qk = tl.dot(q, k, out_dtype=tl.float32)
 
@@ -171,14 +171,14 @@ def sdnq_attn_kernel(
                 v_scale = v_scale_desc.load([start_n])[None, :]
                 p *= v_scale
                 if v.dtype == tl.int8:
-                    p_scale = tl.max(p, 1)[:, None] * (1 / 127.0)
+                    p_scale = tl.mul(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)
+                    p = tl.floor(tl.fma(p, (1 / p_scale), 0.5)).to(tl.int8)
                     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.mul(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)
+                    p = tl.mul(p, (1 / p_scale)).to(v.dtype)
                     acc = tl.fma(tl.dot(p, v, out_dtype=tl.float32), p_scale, acc)
             else:
                 p = p.to(v.dtype)
@@ -201,9 +201,9 @@ def quantize_attn(
     matmul_dtype: str = "int8",
     pv_matmul_dtype: str | None = None,
 ) -> tuple[torch.Tensor]:
-    if matmul_dtype in {"auto", "uint8"}:
+    if matmul_dtype in {"auto", "enabled", "uint8"}:
         matmul_dtype = "int8"
-    if pv_matmul_dtype == "uint8":
+    if pv_matmul_dtype in {"enabled", "uint8"}:
         pv_matmul_dtype = "int8"
     if scale is None:
         scale = q.shape[-1] ** -0.5
@@ -213,7 +213,7 @@ def quantize_attn(
             k = k.sub_(k.mean(dim=2, keepdim=True))
         else:
             k = k.sub(k.mean(dim=2, keepdim=True))
-    if matmul_dtype not in {None, "none", "no"}:
+    if matmul_dtype not in {None, "none", "no", "disabled"}:
         if hadamard is not None:
             q, use_hadamard, hadamard_group_size = apply_hadamard(q, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name="Linear")
             if use_hadamard:
@@ -228,7 +228,7 @@ def quantize_attn(
         k_q = k.contiguous().to(dtype=q.dtype)
         q_scale = None
         k_scale = None
-    if pv_matmul_dtype not in {None, "auto", "none", "no"}:
+    if pv_matmul_dtype not in {None, "auto", "none", "no", "disabled"}:
         quantize_mm_func_pv = quantize_int_mm if pv_matmul_dtype.startswith("int") else quantize_fp_mm
         v_q, v_scale = quantize_mm_func_pv(v.contiguous().to(dtype=torch.float32), dim=-1, matmul_dtype=pv_matmul_dtype)
         v_scale = v_scale.squeeze(-1)
@@ -280,8 +280,8 @@ def get_attn_inputs(
         smooth_k=smooth_k,
         hadamard=hadamard,
         hadamard_group_size=hadamard_group_size,
-        matmul_dtype=matmul_dtype if do_quantize else "no",
-        pv_matmul_dtype=pv_matmul_dtype if do_quantize else "no",
+        matmul_dtype=matmul_dtype if do_quantize else "disabled",
+        pv_matmul_dtype=pv_matmul_dtype if do_quantize else "disabled",
     )
     return query, query_scale, key, key_scale, value, value_scale, attn_mask, scale, out_dtype
 
@@ -308,7 +308,7 @@ def sdnq_triton_atten(
     _, _, VN, VHD = value.shape
 
     hadamard = None
-    if use_hadamard and do_quantize and matmul_dtype not in {None, "none", "no"}:
+    if use_hadamard and do_quantize and matmul_dtype not in {None, "none", "no", "disabled"}:
         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)

modules/sdnq/kernels/triton_mm.py

@@ -14,7 +14,7 @@ matmul_configs = [
     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(",")]
+    for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
 ]
 
 
@@ -38,9 +38,9 @@ def sdnq_triton_mm_kernel(
     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
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = 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)
@@ -72,7 +72,7 @@ def sdnq_triton_mm_kernel(
         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_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 _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = a_desc.load([off_m, off_k])

modules/sdnq/kernels/triton_scaled_mm.py

@@ -1,8 +1,3 @@
-"""
-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
@@ -19,7 +14,7 @@ matmul_configs = [
     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(",")]
+    for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
 ]
 
 
@@ -44,9 +39,9 @@ def sdnq_scaled_mm_kernel(
     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
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = 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)
@@ -78,7 +73,7 @@ def sdnq_scaled_mm_kernel(
         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_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 _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = a_desc.load([off_m, off_k])
@@ -96,17 +91,17 @@ def sdnq_scaled_mm_kernel(
     scale_b = scale_b_desc.load([off_n])[None, :].to(tl.float32)
 
     if bias_ndim == 1:
-        accumulator = accumulator.to(tl.float32) * scale_a
+        accumulator = tl.mul(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
+        accumulator = tl.mul(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 = tl.mul(tl.mul(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))

modules/sdnq/layers/conv/conv_fp16.py

@@ -48,7 +48,7 @@ def conv_fp16_matmul(
             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)
+    input, weight = check_mats(input, weight, matmul_dtype="float16")
 
     if groups == 1:
         result = fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
@@ -70,6 +70,7 @@ def conv_fp16_matmul(
         result = result.permute(0,3,1,2)
     elif conv_type == 3:
         result = result.permute(0,4,1,2,3)
+    result = result.contiguous()
     return result
 
 

modules/sdnq/layers/conv/conv_fp8.py

@@ -46,7 +46,7 @@ def conv_fp8_matmul(
             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)
-    input, weight = check_mats(input, weight)
+    input, weight = check_mats(input, weight, matmul_dtype="float8_e4m3fn")
 
     if groups == 1:
         result = fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
@@ -68,6 +68,7 @@ def conv_fp8_matmul(
         result = result.permute(0,3,1,2)
     elif conv_type == 3:
         result = result.permute(0,4,1,2,3)
+    result = result.contiguous()
     return result
 
 

modules/sdnq/layers/conv/conv_int8.py

@@ -62,7 +62,7 @@ def conv_int8_matmul(
         if bias is not None:
             zero_bias.add_(bias)
         bias = zero_bias
-    input, weight = check_mats(input, weight)
+    input, weight = check_mats(input, weight, matmul_dtype="int8")
 
     if groups == 1:
         result = int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
@@ -84,6 +84,7 @@ def conv_int8_matmul(
         result = result.permute(0,3,1,2)
     elif conv_type == 3:
         result = result.permute(0,4,1,2,3)
+    result = result.contiguous()
     return result
 
 

modules/sdnq/layers/conv/conv_uint8.py

@@ -65,7 +65,7 @@ def conv_uint8_matmul(
         zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point)
     if bias is not None:
         zero_bias.add_(bias)
-    input, weight = check_mats(input, weight)
+    input, weight = check_mats(input, weight, matmul_dtype="uint8")
 
     if groups == 1:
         result = int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(mm_output_shape)
@@ -84,6 +84,7 @@ def conv_uint8_matmul(
         result = result.permute(0,3,1,2)
     elif conv_type == 3:
         result = result.permute(0,4,1,2,3)
+    result = result.contiguous()
     return result
 
 

modules/sdnq/layers/linear/forward.py

@@ -2,12 +2,17 @@
 
 import torch
 
-from ...kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
+from ...kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_contiguous_fp8_mm
 
 
-def check_mats(input: torch.Tensor, weight: torch.Tensor, allow_contiguous_mm: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
-    input = input.contiguous()
-    if allow_contiguous_mm and ((use_contiguous_int8_mm and weight.dtype == torch.int8) or (use_contiguous_fp16_mm and weight.dtype == torch.float16)):
+def check_mats(input: torch.Tensor, weight: torch.Tensor, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.Tensor]:
+    if input is not None:
+        input = input.contiguous()
+    if (
+        (use_contiguous_int8_mm and matmul_dtype in {"int8", "uint8"})
+        or (use_contiguous_fp16_mm and matmul_dtype in {"fp16", "float16"})
+        or (use_contiguous_fp8_mm and matmul_dtype in {"fp8", "float8_e4m3fn"})
+    ):
         weight = weight.contiguous()
     elif weight.is_contiguous():
         weight = weight.t().contiguous().t()

modules/sdnq/layers/linear/linear_fp16.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...common import compile_func
-from ...kernel_wrappers import fp_scaled_mm_func
+from ...kernel_wrappers import fp_scaled_mm_func, include_mm_kernel_in_compile
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
 
@@ -11,7 +11,7 @@ from .forward import check_mats
 from .linear_fp8 import quantize_fp_mm_input
 
 
-def fp16_matmul(
+def get_fp16_matmul_inputs(
     input: torch.FloatTensor,
     weight: torch.Tensor,
     scale: torch.FloatTensor,
@@ -40,7 +40,30 @@ def fp16_matmul(
             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)
+    input, weight = check_mats(input, weight, matmul_dtype="float16")
+    return input, weight, input_scale, scale, bias, return_dtype, output_shape
+
+
+def fp16_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:
+    input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp16_matmul_inputs(
+        input, weight, scale,
+        bias=bias,
+        svd_up=svd_up,
+        svd_down=svd_down,
+        hadamard=hadamard,
+        quantized_weight_shape=quantized_weight_shape,
+        weights_dtype=weights_dtype,
+    )
     return fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
 
 
@@ -69,4 +92,7 @@ def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torc
     )
 
 
-fp16_matmul = compile_func(fp16_matmul)
+if not include_mm_kernel_in_compile:
+    get_fp16_matmul_inputs = compile_func(get_fp16_matmul_inputs)
+else:
+    fp16_matmul = compile_func(fp16_matmul)

modules/sdnq/layers/linear/linear_fp8.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...common import compile_func
-from ...kernel_wrappers import fp8_scaled_mm_func
+from ...kernel_wrappers import fp8_scaled_mm_func, is_fp8_mm_supported, include_mm_kernel_in_compile
 from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
 
@@ -20,7 +20,7 @@ def quantize_fp_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = N
     return input, input_scale
 
 
-def fp8_matmul(
+def get_fp8_matmul_inputs(
     input: torch.FloatTensor,
     weight: torch.Tensor,
     scale: torch.FloatTensor,
@@ -47,7 +47,30 @@ def fp8_matmul(
             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)
-    input, weight = check_mats(input, weight, allow_contiguous_mm=False)
+    input, weight = check_mats(input, weight, matmul_dtype="float8_e4m3fn")
+    return input, weight, input_scale, scale, bias, return_dtype, output_shape
+
+
+def fp8_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:
+    input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp8_matmul_inputs(
+        input, weight, scale,
+        bias=bias,
+        svd_up=svd_up,
+        svd_down=svd_down,
+        hadamard=hadamard,
+        quantized_weight_shape=quantized_weight_shape,
+        weights_dtype=weights_dtype,
+    )
     return fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
 
 
@@ -76,4 +99,7 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch
     )
 
 
-fp8_matmul = compile_func(fp8_matmul)
+if is_fp8_mm_supported and not include_mm_kernel_in_compile:
+    get_fp8_matmul_inputs = compile_func(get_fp8_matmul_inputs)
+else:
+    fp8_matmul = compile_func(fp8_matmul)

modules/sdnq/layers/linear/linear_int8.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...common import compile_func
-from ...kernel_wrappers import int_scaled_mm_func
+from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile
 from ...quant_utils import quantize_int_mm, rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
 
@@ -20,7 +20,7 @@ def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None =
     return input, input_scale
 
 
-def int8_matmul(
+def get_int8_matmul_inputs(
     input: torch.FloatTensor,
     weight: torch.Tensor,
     scale: torch.FloatTensor,
@@ -60,12 +60,36 @@ def int8_matmul(
 
     input, input_scale = quantize_int_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 = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point)
         if bias is not None:
             zero_bias.add_(bias)
         bias = zero_bias
+    input, weight = check_mats(input, weight, matmul_dtype="int8")
+    return input, weight, input_scale, scale, bias, return_dtype, output_shape
 
-    input, weight = check_mats(input, weight)
+
+def int8_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,
+    zero_point: torch.FloatTensor | None = None,
+    hadamard: torch.FloatTensor | None = None,
+    quantized_weight_shape: torch.Size | None = None,
+    weights_dtype: str | None = None,
+) -> torch.FloatTensor:
+    input, weight, input_scale, scale, bias, return_dtype, output_shape = get_int8_matmul_inputs(
+        input, weight, scale,
+        bias=bias,
+        svd_up=svd_up,
+        svd_down=svd_down,
+        zero_point=zero_point,
+        hadamard=hadamard,
+        quantized_weight_shape=quantized_weight_shape,
+        weights_dtype=weights_dtype,
+    )
     return int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
 
 
@@ -96,4 +120,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
     )
 
 
-int8_matmul = compile_func(int8_matmul)
+if not include_mm_kernel_in_compile:
+    get_int8_matmul_inputs = compile_func(get_int8_matmul_inputs)
+else:
+    int8_matmul = compile_func(int8_matmul)

modules/sdnq/layers/linear/linear_uint8.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...common import compile_func
-from ...kernel_wrappers import int_scaled_mm_func
+from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile
 from ...quant_utils import quantize_uint_mm, rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
 
@@ -21,7 +21,7 @@ def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None =
     return input, input_scale, input_zero_point
 
 
-def uint8_matmul(
+def get_uint8_matmul_inputs(
     input: torch.FloatTensor,
     weight: torch.Tensor,
     scale: torch.FloatTensor,
@@ -63,13 +63,38 @@ def uint8_matmul(
     if zero_point is not None:
         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))
+        zero_bias.add_(torch.mul(input_zero_point, zero_point), alpha=input.shape[-1])
     else:
         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)
+    input, weight = check_mats(input, weight, matmul_dtype="uint8")
+    return input, weight, input_scale, scale, zero_bias, return_dtype, output_shape
+
+
+def uint8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    scale: torch.FloatTensor,
+    zero_point: 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:
+    input, weight, input_scale, scale, zero_bias, return_dtype, output_shape = get_uint8_matmul_inputs(
+        input, weight,
+        scale, zero_point,
+        bias=bias,
+        svd_up=svd_up,
+        svd_down=svd_down,
+        hadamard=hadamard,
+        quantized_weight_shape=quantized_weight_shape,
+        weights_dtype=weights_dtype,
+    )
     return int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(output_shape)
 
 
@@ -99,4 +124,7 @@ def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> tor
     )
 
 
-uint8_matmul = compile_func(uint8_matmul)
+if not include_mm_kernel_in_compile:
+    get_uint8_matmul_inputs = compile_func(get_uint8_matmul_inputs)
+else:
+    uint8_matmul = compile_func(uint8_matmul)

modules/sdnq/quant_utils.py

@@ -3,8 +3,8 @@
 import torch
 
 from modules import devices
-from .common import dtype_dict, conv_types, conv_transpose_types
-from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
+from .common import dtype_dict, compile_func, conv_types, conv_transpose_types
+from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_contiguous_fp8_mm
 from .utils import is_pow2, is_pow4, next_power_of_2
 
 
@@ -170,7 +170,11 @@ def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.
 
 @devices.inference_context()
 def prepare_weight_for_matmul(weight: torch.Tensor, matmul_dtype: str | None = "int8") -> torch.Tensor:
-    if (use_contiguous_int8_mm and matmul_dtype in {"int8", "uint8"}) or (use_contiguous_fp16_mm and matmul_dtype == "float16"):
+    if (
+        (use_contiguous_int8_mm and matmul_dtype in {"int8", "uint8"})
+        or (use_contiguous_fp16_mm and matmul_dtype in {"fp16", "float16"})
+        or (use_contiguous_fp8_mm and matmul_dtype in {"fp8", "float8_e4m3fn"})
+    ):
         weight = weight.contiguous()
     elif weight.is_contiguous():
         weight = weight.t_().contiguous().t_()
@@ -225,3 +229,6 @@ def quantize_fp_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.Flo
         weight = weight.add_(torch.randint_like(weight, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32)
     weight = torch.div(weight, scale).nan_to_num_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"])
     return weight, scale
+
+
+rotate_hadamard_compiled = compile_func(rotate_hadamard)

modules/seedvr/config_3b.yaml

@@ -54,7 +54,7 @@ vae:
         - "modules.seedvr.src.models.video_vae_v3.modules.attn_video_vae"
       name: "VideoAutoencoderKLWrapper"
       args: "as_params"
-    freeze_encoder: False
+    freeze_encoder: True
     gradient_checkpoint: True # Disabled to prevent VRAM leaks in inference
   slicing:
     split_size: 4

modules/seedvr/config_7b.yaml

@@ -51,7 +51,7 @@ vae:
         - "modules.seedvr.src.models.video_vae_v3.modules.attn_video_vae"
       name: "VideoAutoencoderKLWrapper"
       args: "as_params"
-    freeze_encoder: False
+    freeze_encoder: True
     # gradient_checkpoint: True
   slicing:
     split_size: 4

modules/seedvr/src/common/cache.py

@@ -21,6 +21,9 @@ class Cache:
             self.cache[key] = result
         return result
 
+    def clear(self):
+        self.cache.clear()
+
     def namespace(self, namespace: str):
         return Cache(
             disable=self.disable,
@@ -31,3 +34,15 @@ class Cache:
     def get(self, key: str):
         key = self.prefix + key
         return self.cache[key]
+
+    def size(self):
+        num = len(self.cache)
+        total_size = 0
+        for value in self.cache.values():
+            if hasattr(value, "element_size") and hasattr(value, "nelement"):
+                total_size += value.element_size() * value.nelement()
+            elif isinstance(value, (list, tuple)):
+                for item in value:
+                    if hasattr(item, "element_size") and hasattr(item, "nelement"):
+                        total_size += item.element_size() * item.nelement()
+        return num, total_size

modules/seedvr/src/common/diffusion/samplers/euler.py

@@ -18,16 +18,10 @@ Euler ODE solver.
 """
 
 from typing import Callable
+import itertools
 import torch
-from einops import rearrange
-from torch.nn import functional as F
-
-#from ....models.dit_v2 import na
-
-from ..types import PredictionType
 from ..utils import expand_dims
 from .base import Sampler, SamplerModelArgs
-import itertools
 
 
 class EulerSampler(Sampler):

modules/seedvr/src/core/generation.py

@@ -84,6 +84,7 @@ def generation_step(runner, text_embeds_dict, cond_latents, temporal_overlap, de
 
     # Process samples with advanced optimization
     samples = optimized_video_rearrange(video_tensors)
+    del video_tensors
     noises = noises[0].to("cpu")
     aug_noises = aug_noises[0].to("cpu")
     cond_latents = cond_latents[0].to("cpu")
@@ -106,7 +107,7 @@ def cut_videos(videos):
     return result
 
 
-def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu'):
+def generation_loop(runner, images, cfg_scale=1.0, cfg_rescale=0.0, steps=1, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu', color_reconstruct=True):
     """
     Main generation loop with context-aware temporal processing
 
@@ -137,9 +138,9 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si
 
     # Configure classifier-free guidance
     runner.config.diffusion.cfg.scale = cfg_scale
-    runner.config.diffusion.cfg.rescale = 0.0
+    runner.config.diffusion.cfg.rescale = cfg_rescale
     # Configure sampling steps
-    runner.config.diffusion.timesteps.sampling.steps = 1
+    runner.config.diffusion.timesteps.sampling.steps = steps
     runner.configure_diffusion()
 
     # Set random seed
@@ -159,7 +160,8 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si
     ])
 
     # Initialize generation state
-    batch_samples = []
+    final_video_images = None
+    current_idx = 0
 
     # Load text embeddings with adaptive dtype
     text_embeds = {"texts_pos": [runner.text_pos_embeds], "texts_neg": [runner.text_neg_embeds]}
@@ -215,7 +217,8 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si
 
         # Normal generation
         samples = generation_step(runner, text_embeds, cond_latents=cond_latents, temporal_overlap=temporal_overlap, device=device)
-        #del cond_latents
+        if samples is None:
+            return
         del cond_latents
 
         # Post-process samples
@@ -223,47 +226,35 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si
         del samples
         #del samples
         if ori_lengths[0] < sample.shape[0]:
-            sample = sample[:ori_lengths[0]]
+            sample = sample[:ori_lengths[0]].contiguous()
 
         # Apply color correction if available
-        transformed_video = transformed_video.to(device)
-        input_video = [optimized_single_video_rearrange(transformed_video)]
-        del transformed_video
-        sample = wavelet_reconstruction(sample, input_video[0][:sample.size(0)])
-        del input_video
+        if color_reconstruct:
+            transformed_video = transformed_video.to(device)
+            input_video = [optimized_single_video_rearrange(transformed_video)]
+            del transformed_video
+            sample = wavelet_reconstruction(sample, input_video[0][:sample.size(0)])
+            del input_video
 
         # Convert to final image format
         sample = optimized_sample_to_image_format(sample)
         sample = sample.clip(-1, 1).mul_(0.5).add_(0.5)
-        sample_cpu = sample.to(torch.float16).to("cpu")
+        sample = sample.detach().to(torch.float16, non_blocking=True).cpu()
+        if final_video_images is None:
+            total_frames = len(images)
+            H, W, C = sample.shape[1], sample.shape[2], sample.shape[3]
+            final_video_images = torch.empty((total_frames, H, W, C), dtype=torch.float16)
+
+        batch_frames = sample.shape[0]
+        final_video_images[current_idx:current_idx + batch_frames] = sample
+        current_idx += batch_frames
         del sample
-        batch_samples.append(sample_cpu)
-        #del sample
 
         if progress_callback:
             progress_callback(batch_count+1, total_batches, current_frames, "Processing batch...")
 
 
-    # 1. Calculer la taille totale finale
-    total_frames = sum(batch.shape[0] for batch in batch_samples)
-    if len(batch_samples) > 0:
-        sample_shape = batch_samples[0].shape
-        H, W, C = sample_shape[1], sample_shape[2], sample_shape[3]
-        final_video_images = torch.empty((total_frames, H, W, C), dtype=torch.float16)
-        block_size = 500
-        current_idx = 0
-
-        for block_start in range(0, len(batch_samples), block_size):
-            block_end = min(block_start + block_size, len(batch_samples))
-            current_block = []
-            for i in range(block_start, block_end):
-                current_block.append(batch_samples[i].to(device))
-            block_result = torch.cat(current_block, dim=0)
-            block_frames = block_result.shape[0]
-            final_video_images[current_idx:current_idx + block_frames] = block_result.to("cpu")
-            current_idx += block_frames
-            del current_block, block_result
-    else:
+    if final_video_images is None:
         print("SeedVR2: No batch_samples to process")
         final_video_images = torch.empty((0, 0, 0, 0), dtype=torch.float16)
 

modules/seedvr/src/core/infer.py

@@ -7,6 +7,7 @@ from modules.seedvr.src.models.dit_v2 import na
 
 if TYPE_CHECKING:
     from modules.seedvr.src.models.dit_v2.nadit import NaDiT
+    from modules.seedvr.src.models.video_vae_v3.modules.attn_video_vae import VideoAutoencoderKLWrapper
 
 
 def optimized_channels_to_last(tensor: torch.Tensor) -> torch.Tensor:
@@ -49,7 +50,7 @@ class SeedVRPipeline():
         self.config = config
         self.device = device
         self.dtype = dtype
-        self.vae = None
+        self.vae: VideoAutoencoderKLWrapper = None
         self.dit: NaDiT = None
         self.sampler = None
         self.schedule = None
@@ -132,6 +133,7 @@ class SeedVRPipeline():
                 latent = rearrange(latent, "b c ... -> b ... c")
                 #latent = optimized_channels_to_last(latent)
                 latent = (latent - shift) * scale
+                latent = latent.contiguous()
                 latents.append(latent)
 
             # Ungroup back to individual latent with the original order.
@@ -174,10 +176,11 @@ class SeedVRPipeline():
                 latent = latent / scale + shift
                 latent = rearrange(latent, "b ... c -> b c ...")
                 #latent = optimized_channels_to_second(latent)
-                latent = latent.squeeze(2)
+                latent = latent.squeeze(2).contiguous()
 
                 # 🚀 OPTIMISATION 3: Décodage direct SANS autocast (utilise l'autocast externe)
                 sample = self.vae.decode(latent).sample
+                del latent
                 #sample = self.vae.decode(latent).sample
                 #sample = self.vae.decode(latent).sample
 
@@ -186,6 +189,7 @@ class SeedVRPipeline():
                     sample = self.vae.postprocess(sample)
 
                 samples.append(sample)
+                del sample
 
             # Ungroup back to individual sample with the original order.
             if self.config.vae.grouping:
@@ -277,9 +281,9 @@ class SeedVRPipeline():
             text_neg_embeds, text_neg_shapes = na.flatten(texts_neg)
 
         # Adapter les embeddings texte au dtype cible (compatible avec FP8)
-        if isinstance(text_pos_embeds, torch.Tensor):
+        if isinstance(text_pos_embeds, torch.Tensor) and text_pos_embeds.dtype != target_dtype:
             text_pos_embeds = text_pos_embeds.to(target_dtype)
-        if isinstance(text_neg_embeds, torch.Tensor):
+        if isinstance(text_neg_embeds, torch.Tensor) and text_neg_embeds.dtype != target_dtype:
             text_neg_embeds = text_neg_embeds.to(target_dtype)
 
         # Flatten.
@@ -289,7 +293,9 @@ class SeedVRPipeline():
         # Adapter les latents au dtype cible (compatible avec FP8)
         latents = latents.to(target_dtype) if latents.dtype != target_dtype else latents
         latents_cond = latents_cond.to(target_dtype) if latents_cond.dtype != target_dtype else latents_cond
-        self.dit = self.dit.to(device=self.device, dtype=target_dtype)
+        current_dit_param = next(self.dit.parameters())
+        if current_dit_param.dtype != target_dtype or current_dit_param.device != torch.device(self.device):
+            self.dit = self.dit.to(device=self.device, dtype=target_dtype)
 
         latents = self.sampler.sample(
             x=latents,
@@ -309,10 +315,7 @@ class SeedVRPipeline():
                     timestep=args.t.repeat(batch_size),
                 ).vid_sample,
                 scale=(
-                    cfg_scale
-                    if (args.i + 1) / len(self.sampler.timesteps)
-                    <= self.config.diffusion.cfg.get("partial", 1)
-                    else 1.0
+                    cfg_scale if (args.i + 1) / len(self.sampler.timesteps) <= self.config.diffusion.cfg.get("partial", 1) else 1.0
                 ),
                 rescale=self.config.diffusion.cfg.rescale,
             ),
@@ -324,6 +327,7 @@ class SeedVRPipeline():
         vae_dtype = self.vae.dtype
         decode_dtype = torch.float16 if (vae_dtype == torch.float16 or target_dtype == torch.float16) else vae_dtype
         samples = self.vae_decode(latents, target_dtype=decode_dtype)
+        del latents
 
         if samples and len(samples) > 0 and samples[0].dtype != torch.float16:
             samples = [sample.to(torch.float16, non_blocking=True) for sample in samples]

modules/seedvr/src/data/image/transforms/side_resize.py

@@ -50,5 +50,6 @@ class SideResize:
             size = min(width, height)
         else:
             size = self.size
-
-        return TVF.resize(image, size, self.interpolation)
+        size_w = int(size) // 8 * 8
+        size_h = int(size * height / width) // 8 * 8
+        return TVF.resize(image, (size_h, size_w), self.interpolation)

modules/seedvr/src/models/dit/patch.py

@@ -42,6 +42,8 @@ class PatchIn(nn.Module):
     ) -> torch.Tensor:
         t, h, w = self.patch_size
         vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         return vid
 
@@ -63,30 +65,41 @@ class PatchOut(nn.Module):
         vid: torch.Tensor,
     ) -> torch.Tensor:
         t, h, w = self.patch_size
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
         return vid
 
 
 class NaPatchIn(PatchIn):
-    def forward(
+    def forward( # pylint: disable=arguments-differ
         self,
         vid: torch.Tensor,  # l c
         vid_shape: torch.LongTensor,
     ) -> torch.Tensor:
         t, h, w = self.patch_size
-        if not (t == h == w == 1):
-            vid, vid_shape = na.rearrange(
-                vid, vid_shape, "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w
-            )
+        if not t == h == w == 1:
+            vid = na.unflatten(vid, vid_shape)
+            for i in range(len(vid)):
+                if t > 1 and vid_shape[i, 0] % t != 0:
+                    vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0)
+                if h > 1 and vid_shape[i, 1] % h != 0:
+                    vid[i] = torch.cat([vid[i][:, :1]] * (h - vid[i].size(1) % h) + [vid[i]], dim=1)
+                if w > 1 and vid_shape[i, 2] % w != 0:
+                    vid[i] = torch.cat([vid[i][:, :, :1]] * (w - vid[i].size(2) % w) + [vid[i]], dim=2)
+                vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w)
+            vid, vid_shape = na.flatten(vid)
         # slice vid after patching in when using sequence parallelism
         vid = slice_inputs(vid, dim=0)
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         return vid, vid_shape
 
 
 class NaPatchOut(PatchOut):
-    def forward(
+    def forward( # pylint: disable=arguments-differ
         self,
         vid: torch.FloatTensor,  # l c
         vid_shape: torch.LongTensor,
@@ -96,8 +109,10 @@ class NaPatchOut(PatchOut):
         torch.LongTensor,
     ]:
         t, h, w = self.patch_size
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
-        # gather vid before patching out when enabling sequence parallelism
+        # gather vid before patchting out when enabling sequence parallelism
         vid = gather_outputs(
             vid,
             gather_dim=0,
@@ -105,8 +120,15 @@ class NaPatchOut(PatchOut):
             unpad_shape=vid_shape,
             cache=cache.namespace("vid"),
         )
-        if not (t == h == w == 1):
-            vid, vid_shape = na.rearrange(
-                vid, vid_shape, "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w
-            )
+        if not t == h == w == 1:
+            vid = na.unflatten(vid, vid_shape)
+            for i in range(len(vid)):
+                vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w)
+                if t > 1 and vid_shape[i, 0] % t != 0:
+                    vid[i] = vid[i][(t - vid_shape[i, 0] % t) :]
+                if h > 1 and vid_shape[i, 1] % h != 0:
+                    vid[i] = vid[i][:, (h - vid_shape[i, 1] % h) :]
+                if w > 1 and vid_shape[i, 2] % w != 0:
+                    vid[i] = vid[i][:, :, (w - vid_shape[i, 2] % w) :]
+            vid, vid_shape = na.flatten(vid)
         return vid, vid_shape

modules/seedvr/src/models/dit_v2/nablocks/attention/mmattn.py

@@ -13,12 +13,11 @@
 # // limitations under the License.
 
 from typing import Optional, Tuple, Union
+from itertools import chain
 import torch
 from einops import rearrange
 from torch import nn
-from torch.nn import functional as F
 from torch.nn.modules.utils import _triple
-
 from .....common.cache import Cache
 from .....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
 from .....common.half_precision_fixes import safe_pad_operation
@@ -29,7 +28,6 @@ from ...mm import MMArg, MMModule
 from ...normalization import norm_layer_type
 from ...rope import get_na_rope
 from ...window import get_window_op
-from itertools import chain
 
 
 class NaMMAttention(nn.Module):

modules/seedvr/src/models/dit_v2/patch/patch_v1.py

@@ -45,6 +45,8 @@ class PatchIn(nn.Module):
             assert vid.size(2) % t == 1
             vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2)
         vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         return vid
 
@@ -83,16 +85,22 @@ class NaPatchIn(PatchIn):
         cache = cache.namespace("patch")
         vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape)
         t, h, w = self.patch_size
-        if not (t == h == w == 1):
+        if not t == h == w == 1:
             vid = na.unflatten(vid, vid_shape)
             for i in range(len(vid)):
                 if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
                     vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0)
+                if h > 1 and vid_shape_before_patchify[i, 1] % h != 0:
+                    vid[i] = torch.cat([vid[i][:, :1]] * (h - vid[i].size(1) % h) + [vid[i]], dim=1)
+                if w > 1 and vid_shape_before_patchify[i, 2] % w != 0:
+                    vid[i] = torch.cat([vid[i][:, :, :1]] * (w - vid[i].size(2) % w) + [vid[i]], dim=2)
                 vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w)
             vid, vid_shape = na.flatten(vid)
 
         # slice vid after patching in when using sequence parallelism
         vid = slice_inputs(vid, dim=0)
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         return vid, vid_shape
 
@@ -111,17 +119,23 @@ class NaPatchOut(PatchOut):
         vid_shape_before_patchify = cache.get("vid_shape_before_patchify")
 
         t, h, w = self.patch_size
+        if vid.dtype != self.proj.weight.dtype:
+            vid = vid.to(self.proj.weight.dtype)
         vid = self.proj(vid)
         # gather vid before patching out when enabling sequence parallelism
         vid = gather_outputs(
             vid, gather_dim=0, padding_dim=0, unpad_shape=vid_shape, cache=cache.namespace("vid")
         )
-        if not (t == h == w == 1):
+        if not t == h == w == 1:
             vid = na.unflatten(vid, vid_shape)
             for i in range(len(vid)):
                 vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w)
                 if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
                     vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :]
+                if h > 1 and vid_shape_before_patchify[i, 1] % h != 0:
+                    vid[i] = vid[i][:, (h - vid_shape_before_patchify[i, 1] % h) :]
+                if w > 1 and vid_shape_before_patchify[i, 2] % w != 0:
+                    vid[i] = vid[i][:, :, (w - vid_shape_before_patchify[i, 2] % w) :]
             vid, vid_shape = na.flatten(vid)
 
         return vid, vid_shape

modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py

@@ -114,14 +114,26 @@ class Upsample3D(Upsample2D):
             hidden_states = [hidden_states]
         # ADD BY NUMZ
         for i in range(len(hidden_states)):
-            hidden_states[i] = self.upscale_conv(hidden_states[i])
-            hidden_states[i] = rearrange(
-                hidden_states[i],
-                "b (x y z c) f h w -> b c (f z) (h x) (w y)",
-                x=self.spatial_ratio,
-                y=self.spatial_ratio,
-                z=self.temporal_ratio,
-            )
+            if self.use_conv and hasattr(self, "upscale_conv") and self.upscale_conv.kernel_size == (1, 1, 1):
+                hidden_states[i] = hidden_states[i].repeat_interleave(self.temporal_ratio, dim=2)
+                if self.spatial_ratio != 1:
+                    hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=3)
+                    hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=4)
+            elif self.use_conv:
+                hidden_states[i] = self.upscale_conv(hidden_states[i])
+                hidden_states[i] = rearrange(
+                    hidden_states[i],
+                    "b (x y z c) f h w -> b c (f z) (h x) (w y)",
+                    x=self.spatial_ratio,
+                    y=self.spatial_ratio,
+                    z=self.temporal_ratio,
+                ).contiguous()
+            else:
+                if self.temporal_ratio != 1:
+                    hidden_states[i] = hidden_states[i].repeat_interleave(self.temporal_ratio, dim=2)
+                if self.spatial_ratio != 1:
+                    hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=3)
+                    hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=4)
 
         # [Overridden] For causal temporal conv
         if self.temporal_up and memory_state != MemoryState.ACTIVE:
@@ -1044,7 +1056,7 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         norm_num_groups: int = 32,
         sample_size: int = 32,
         scaling_factor: float = 0.18215,
-        force_upcast: float = True,
+        force_upcast: float = False,
         attention: bool = True,
         temporal_scale_num: int = 0,
         slicing_up_num: int = 0,
@@ -1146,9 +1158,13 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
 
     @apply_forward_hook
     def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:
-        # h = self.slicing_encode(x)
-        h = self.tiled_encode(x)
-        posterior = DiagonalGaussianDistribution(h)
+        if self.use_slicing_encode:
+            encoded = self.slicing_encode(x)
+        elif self.use_tiling_encode:
+            encoded = self.tiled_encode(x)
+        else:
+            encoded = self._encode(x)
+        posterior = DiagonalGaussianDistribution(encoded)
 
         if not return_dict:
             return (posterior,)
@@ -1159,8 +1175,12 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
     def decode(
         self, z: torch.Tensor, return_dict: bool = True
     ) -> Union[DecoderOutput, torch.Tensor]:
-        # decoded = self.slicing_decode(z)
-        decoded = self.tiled_decode(z)
+        if self.use_slicing_decode:
+            decoded = self.slicing_decode(z)
+        elif self.use_tiling_decode:
+            decoded = self.tiled_decode(z)
+        else:
+            decoded = self._decode(z)
 
         if not return_dict:
             return (decoded,)
@@ -1170,8 +1190,7 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
     def _encode(
         self, x: torch.Tensor, memory_state: MemoryState = MemoryState.DISABLED
     ) -> torch.Tensor:
-        _x = x.to(self.device)
-        _x = causal_conv_slice_inputs(_x, self.slicing_sample_min_size, memory_state=memory_state)
+        _x = causal_conv_slice_inputs(x.to(self.device), self.slicing_sample_min_size, memory_state=memory_state)
         h = self.encoder(_x, memory_state=memory_state)
         if self.quant_conv is not None:
             output = self.quant_conv(h, memory_state=memory_state)
@@ -1243,51 +1262,109 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
         blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
         row_limit = self.tile_latent_min_size - blend_extent
-        rows = []
-        for i in range(0, x.shape[3], overlap_size):
-            row = []
-            for j in range(0, x.shape[4], overlap_size):
+        prev_row = None
+        self.tiles = 0
+
+        row_positions = list(range(0, x.shape[3], overlap_size))
+        col_positions = list(range(0, x.shape[4], overlap_size))
+        enc = None
+        output_width = 0
+        h_cursor = 0
+
+        for _row_idx, i in enumerate(row_positions):
+            row_tiles = []
+            for tile_idx, j in enumerate(col_positions):
                 tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                 tile = self._encode(tile)
-                row.append(tile)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=4))
-        enc = torch.cat(result_rows, dim=3)
-        return enc
+                if tile.ndim == 4:
+                    tile = tile.unsqueeze(0)
+                if prev_row is not None:
+                    tile = self.blend_v(prev_row[tile_idx], tile, blend_extent)
+                if tile_idx > 0:
+                    tile = self.blend_h(row_tiles[-1], tile, blend_extent)
+                row_tiles.append(tile)
+                self.tiles += 1
+
+            cropped_tiles = [tile[:, :, :, :row_limit, :row_limit] for tile in row_tiles]
+            row_width = 0
+            for cropped in cropped_tiles:
+                row_width += cropped.shape[-1]
+                if output_width < cropped.shape[-1]:
+                    output_width = cropped.shape[-1]
+
+            if enc is None:
+                enc = torch.empty(
+                    cropped_tiles[0].shape[0],
+                    cropped_tiles[0].shape[1],
+                    cropped_tiles[0].shape[2],
+                    len(row_positions) * row_limit,
+                    len(col_positions) * row_limit,
+                    dtype=cropped_tiles[0].dtype,
+                    device=cropped_tiles[0].device,
+                )
+
+            w_cursor = 0
+            for cropped in cropped_tiles:
+                enc[:, :, :, h_cursor : h_cursor + cropped.shape[-2], w_cursor : w_cursor + cropped.shape[-1]] = cropped
+                w_cursor += cropped.shape[-1]
+
+            h_cursor += cropped_tiles[0].shape[-2]
+            prev_row = row_tiles
+
+        return enc[:, :, :, :h_cursor, :w_cursor]
 
     def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
         overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
         blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
         row_limit = self.tile_sample_min_size - blend_extent
-        rows = []
-        for i in range(0, z.shape[3], overlap_size):
-            row = []
-            for j in range(0, z.shape[4], overlap_size):
+        prev_row = None
+
+        row_positions = list(range(0, z.shape[3], overlap_size))
+        col_positions = list(range(0, z.shape[4], overlap_size))
+        dec = None
+        output_width = 0
+        h_cursor = 0
+
+        for _row_idx, i in enumerate(row_positions):
+            row_tiles = []
+            for tile_idx, j in enumerate(col_positions):
                 tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                 decoded = self.decoder(tile)
-                row.append(decoded)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=4))
-        dec = torch.cat(result_rows, dim=3)
-        return dec
+                if decoded.ndim == 4:
+                    decoded = decoded.unsqueeze(0)
+                if prev_row is not None:
+                    decoded = self.blend_v(prev_row[tile_idx], decoded, blend_extent)
+                if tile_idx > 0:
+                    decoded = self.blend_h(row_tiles[-1], decoded, blend_extent)
+                row_tiles.append(decoded)
+
+            cropped_tiles = [tile[:, :, :, :row_limit, :row_limit] for tile in row_tiles]
+            row_width = 0
+            for cropped in cropped_tiles:
+                row_width += cropped.shape[-1]
+                if output_width < cropped.shape[-1]:
+                    output_width = cropped.shape[-1]
+
+            if dec is None:
+                dec = torch.empty(
+                    cropped_tiles[0].shape[0],
+                    cropped_tiles[0].shape[1],
+                    cropped_tiles[0].shape[2],
+                    len(row_positions) * row_limit,
+                    len(col_positions) * row_limit,
+                    dtype=cropped_tiles[0].dtype,
+                    device=cropped_tiles[0].device,
+                )
+
+            w_cursor = 0
+            for cropped in cropped_tiles:
+                dec[:, :, :, h_cursor : h_cursor + cropped.shape[-2], w_cursor : w_cursor + cropped.shape[-1]] = cropped
+                w_cursor += cropped.shape[-1]
+
+            h_cursor += cropped_tiles[0].shape[-2]
+            prev_row = row_tiles
+
+        return dec[:, :, :, :h_cursor, :w_cursor]
 
     def forward(
         self, x: torch.FloatTensor, mode: Literal["encode", "decode", "all"] = "all", **kwargs
@@ -1328,7 +1405,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
     ):
         self.spatial_downsample_factor = spatial_downsample_factor
         self.temporal_downsample_factor = temporal_downsample_factor
-        self.freeze_encoder = freeze_encoder
+        self.freeze_encoder = True
         super().__init__(*args, **kwargs)
 
     def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:

modules/seedvr/src/models/video_vae_v3/modules/causal_inflation_lib.py

@@ -150,6 +150,8 @@ class InflatedCausalConv3d(Conv3d):
         assert memory_state != MemoryState.UNSET
         if memory_state != MemoryState.ACTIVE:
             self.memory = None
+        if torch.is_tensor(input) and memory_state == MemoryState.DISABLED:
+            return self.basic_forward(input, memory_state)
         if (
             math.isinf(self.memory_limit)
             and torch.is_tensor(input)

modules/seedvr/src/models/video_vae_v3/modules/video_vae.py.old

@@ -1,936 +0,0 @@
-# Copyright (c) 2023 HuggingFace Team
-# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
-# SPDX-License-Identifier: Apache License, Version 2.0 (the "License")
-#
-# This file has been modified by ByteDance Ltd. and/or its affiliates. on 1st June 2025
-#
-# Original file was released under Apache License, Version 2.0 (the "License"), with the full license text
-# available at http://www.apache.org/licenses/LICENSE-2.0.
-#
-# This modified file is released under the same license.
-
-from contextlib import nullcontext
-from typing import Optional, Tuple, Literal, Callable, Union
-import torch
-import torch.nn as nn
-from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
-from einops import rearrange
-from ....common.half_precision_fixes import safe_pad_operation
-from ....common.logger import get_logger
-from .causal_inflation_lib import InflatedCausalConv3d, causal_norm_wrapper, init_causal_conv3d, remove_head
-from .context_parallel_lib import causal_conv_gather_outputs, causal_conv_slice_inputs
-from .global_config import set_norm_limit
-from .types import CausalAutoencoderOutput, CausalDecoderOutput, CausalEncoderOutput, MemoryState, _inflation_mode_t, _memory_device_t, _receptive_field_t, _selective_checkpointing_t
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-# Fake func, no checkpointing is required for inference
-def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
-    return module(*args, **kwargs)
-
-class ResnetBlock2D(nn.Module):
-    r"""
-    A Resnet block.
-
-    Parameters:
-        in_channels (`int`): The number of channels in the input.
-        out_channels (`int`, *optional*, default to be `None`):
-            The number of output channels for the first conv2d layer.
-            If None, same as `in_channels`.
-        dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
-    """
-
-    def __init__(
-        self, *, in_channels: int, out_channels: Optional[int] = None, dropout: float = 0.0
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-
-        self.nonlinearity = nn.SiLU()
-
-        self.norm1 = torch.nn.GroupNorm(
-            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-        )
-
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        self.norm2 = torch.nn.GroupNorm(
-            num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
-        )
-
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        self.use_in_shortcut = self.in_channels != out_channels
-
-        self.conv_shortcut = None
-        if self.use_in_shortcut:
-            self.conv_shortcut = nn.Conv2d(
-                in_channels, out_channels, kernel_size=1, stride=1, padding=0
-            )
-
-    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
-        hidden = input_tensor
-
-        hidden = self.norm1(hidden)
-        hidden = self.nonlinearity(hidden)
-        hidden = self.conv1(hidden)
-
-        hidden = self.norm2(hidden)
-        hidden = self.nonlinearity(hidden)
-        hidden = self.dropout(hidden)
-        hidden = self.conv2(hidden)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor)
-
-        output_tensor = input_tensor + hidden
-
-        return output_tensor
-
-class Upsample3D(nn.Module):
-    """A 3D upsampling layer."""
-
-    def __init__(
-        self,
-        channels: int,
-        inflation_mode: _inflation_mode_t = "tail",
-        temporal_up: bool = False,
-        spatial_up: bool = True,
-        slicing: bool = False,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.conv = init_causal_conv3d(
-            self.channels, self.channels, kernel_size=3, padding=1, inflation_mode=inflation_mode
-        )
-
-        self.temporal_up = temporal_up
-        self.spatial_up = spatial_up
-        self.temporal_ratio = 2 if temporal_up else 1
-        self.spatial_ratio = 2 if spatial_up else 1
-        self.slicing = slicing
-
-        upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio
-        self.upscale_conv = nn.Conv3d(
-            self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0
-        )
-        identity = (
-            torch.eye(self.channels).repeat(upscale_ratio, 1).reshape_as(self.upscale_conv.weight)
-        )
-
-        self.upscale_conv.weight.data.copy_(identity)
-        nn.init.zeros_(self.upscale_conv.bias)
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        memory_state: MemoryState,
-    ) -> torch.FloatTensor:
-        return gradient_checkpointing(
-            self.custom_forward,
-            hidden_states,
-            memory_state,
-            enabled=self.training and self.gradient_checkpointing,
-        )
-
-    def custom_forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        memory_state: MemoryState,
-    ) -> torch.FloatTensor:
-        assert hidden_states.shape[1] == self.channels
-
-        if self.slicing:
-            split_size = hidden_states.size(2) // 2
-            hidden_states = list(
-                hidden_states.split([split_size, hidden_states.size(2) - split_size], dim=2)
-            )
-        else:
-            hidden_states = [hidden_states]
-
-        for i in range(len(hidden_states)):
-            hidden_states[i] = self.upscale_conv(hidden_states[i])
-            hidden_states[i] = rearrange(
-                hidden_states[i],
-                "b (x y z c) f h w -> b c (f z) (h x) (w y)",
-                x=self.spatial_ratio,
-                y=self.spatial_ratio,
-                z=self.temporal_ratio,
-            )
-
-        # [Overridden] For causal temporal conv
-        if self.temporal_up and memory_state != MemoryState.ACTIVE:
-            hidden_states[0] = remove_head(hidden_states[0])
-
-        if self.slicing:
-            hidden_states = self.conv(hidden_states, memory_state=memory_state)
-            return torch.cat(hidden_states, dim=2)
-        else:
-            return self.conv(hidden_states[0], memory_state=memory_state)
-
-
-class Downsample3D(nn.Module):
-    """A 3D downsampling layer."""
-
-    def __init__(
-        self,
-        channels: int,
-        inflation_mode: _inflation_mode_t = "tail",
-        temporal_down: bool = False,
-        spatial_down: bool = True,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.temporal_down = temporal_down
-        self.spatial_down = spatial_down
-
-        self.temporal_ratio = 2 if temporal_down else 1
-        self.spatial_ratio = 2 if spatial_down else 1
-
-        self.temporal_kernel = 3 if temporal_down else 1
-        self.spatial_kernel = 3 if spatial_down else 1
-
-        self.conv = init_causal_conv3d(
-            self.channels,
-            self.channels,
-            kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel),
-            stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio),
-            padding=((1 if self.temporal_down else 0), 0, 0),
-            inflation_mode=inflation_mode,
-        )
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        memory_state: MemoryState,
-    ) -> torch.FloatTensor:
-        return gradient_checkpointing(
-            self.custom_forward,
-            hidden_states,
-            memory_state,
-            enabled=self.training and self.gradient_checkpointing,
-        )
-
-    def custom_forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        memory_state: MemoryState,
-    ) -> torch.FloatTensor:
-
-        assert hidden_states.shape[1] == self.channels
-
-        if self.spatial_down:
-            hidden_states = safe_pad_operation(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
-
-        hidden_states = self.conv(hidden_states, memory_state=memory_state)
-        return hidden_states
-
-
-class ResnetBlock3D(ResnetBlock2D):
-    def __init__(
-        self,
-        *args,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-        self.conv1 = init_causal_conv3d(
-            self.in_channels,
-            self.out_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            inflation_mode=inflation_mode,
-        )
-
-        self.conv2 = init_causal_conv3d(
-            self.out_channels,
-            self.out_channels,
-            kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3),
-            stride=1,
-            padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1),
-            inflation_mode=inflation_mode,
-        )
-
-        if self.use_in_shortcut:
-            self.conv_shortcut = init_causal_conv3d(
-                self.in_channels,
-                self.out_channels,
-                kernel_size=1,
-                stride=1,
-                padding=0,
-                bias=(self.conv_shortcut.bias is not None),
-                inflation_mode=inflation_mode,
-            )
-        self.gradient_checkpointing = False
-
-    def forward(self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET):
-        return gradient_checkpointing(
-            self.custom_forward,
-            input_tensor,
-            memory_state,
-            enabled=self.training and self.gradient_checkpointing,
-        )
-
-    def custom_forward(
-        self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET
-    ):
-        assert memory_state != MemoryState.UNSET
-        hidden_states = input_tensor
-
-        hidden_states = causal_norm_wrapper(self.norm1, hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.conv1(hidden_states, memory_state=memory_state)
-
-        hidden_states = causal_norm_wrapper(self.norm2, hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states, memory_state=memory_state)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state)
-
-        output_tensor = input_tensor + hidden_states
-
-        return output_tensor
-
-
-class DownEncoderBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        add_downsample: bool = True,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        temporal_down: bool = True,
-        spatial_down: bool = True,
-    ):
-        super().__init__()
-        resnets = []
-
-        for i in range(num_layers):
-            in_channels = in_channels if i == 0 else out_channels
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=in_channels,
-                    out_channels=out_channels,
-                    dropout=dropout,
-                    inflation_mode=inflation_mode,
-                    time_receptive_field=time_receptive_field,
-                )
-            )
-
-        self.resnets = nn.ModuleList(resnets)
-
-        self.downsamplers = None
-        if add_downsample:
-            # Todo: Refactor this line before V5 Image VAE Training.
-            self.downsamplers = nn.ModuleList(
-                [
-                    Downsample3D(
-                        channels=out_channels,
-                        inflation_mode=inflation_mode,
-                        temporal_down=temporal_down,
-                        spatial_down=spatial_down,
-                    )
-                ]
-            )
-
-    def forward(
-        self, hidden_states: torch.FloatTensor, memory_state: MemoryState
-    ) -> torch.FloatTensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, memory_state=memory_state)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states, memory_state=memory_state)
-
-        return hidden_states
-
-
-class UpDecoderBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        add_upsample: bool = True,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        temporal_up: bool = True,
-        spatial_up: bool = True,
-        slicing: bool = False,
-    ):
-        super().__init__()
-        resnets = []
-
-        for i in range(num_layers):
-            input_channels = in_channels if i == 0 else out_channels
-
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=input_channels,
-                    out_channels=out_channels,
-                    dropout=dropout,
-                    inflation_mode=inflation_mode,
-                    time_receptive_field=time_receptive_field,
-                )
-            )
-
-        self.resnets = nn.ModuleList(resnets)
-
-        self.upsamplers = None
-        # Todo: Refactor this line before V5 Image VAE Training.
-        if add_upsample:
-            self.upsamplers = nn.ModuleList(
-                [
-                    Upsample3D(
-                        channels=out_channels,
-                        inflation_mode=inflation_mode,
-                        temporal_up=temporal_up,
-                        spatial_up=spatial_up,
-                        slicing=slicing,
-                    )
-                ]
-            )
-
-    def forward(
-        self, hidden_states: torch.FloatTensor, memory_state: MemoryState
-    ) -> torch.FloatTensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, memory_state=memory_state)
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, memory_state=memory_state)
-
-        return hidden_states
-
-
-class UNetMidBlock3D(nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        dropout: float = 0.0,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-    ):
-        super().__init__()
-        self.resnets = nn.ModuleList(
-            [
-                ResnetBlock3D(
-                    in_channels=channels,
-                    out_channels=channels,
-                    dropout=dropout,
-                    inflation_mode=inflation_mode,
-                    time_receptive_field=time_receptive_field,
-                ),
-                ResnetBlock3D(
-                    in_channels=channels,
-                    out_channels=channels,
-                    dropout=dropout,
-                    inflation_mode=inflation_mode,
-                    time_receptive_field=time_receptive_field,
-                ),
-            ]
-        )
-
-    def forward(self, hidden_states: torch.Tensor, memory_state: MemoryState):
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, memory_state)
-        return hidden_states
-
-
-class Encoder3D(nn.Module):
-    r"""
-    The `Encoder` layer of a variational autoencoder that encodes
-    its input into a latent representation.
-    """
-
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        block_out_channels: Tuple[int, ...] = (64,),
-        layers_per_block: int = 2,
-        double_z: bool = True,
-        temporal_down_num: int = 2,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
-    ):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-
-        self.temporal_down_num = temporal_down_num
-
-        self.conv_in = init_causal_conv3d(
-            in_channels,
-            block_out_channels[0],
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            inflation_mode=inflation_mode,
-        )
-
-        self.down_blocks = nn.ModuleList([])
-
-        # down
-        output_channel = block_out_channels[0]
-        for i in range(len(block_out_channels)):
-            input_channel = output_channel
-            output_channel = block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-            is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1
-            # Note: take the last one
-
-            down_block = DownEncoderBlock3D(
-                num_layers=self.layers_per_block,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                add_downsample=not is_final_block,
-                temporal_down=is_temporal_down_block,
-                spatial_down=True,
-                inflation_mode=inflation_mode,
-                time_receptive_field=time_receptive_field,
-            )
-            self.down_blocks.append(down_block)
-
-        # mid
-        self.mid_block = UNetMidBlock3D(
-            channels=block_out_channels[-1],
-            inflation_mode=inflation_mode,
-            time_receptive_field=time_receptive_field,
-        )
-
-        # out
-        self.conv_norm_out = nn.GroupNorm(
-            num_channels=block_out_channels[-1], num_groups=32, eps=1e-6
-        )
-        self.conv_act = nn.SiLU()
-
-        conv_out_channels = 2 * out_channels if double_z else out_channels
-        self.conv_out = init_causal_conv3d(
-            block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode
-        )
-
-        assert len(selective_checkpointing) == len(self.down_blocks)
-        self.set_gradient_checkpointing(selective_checkpointing)
-
-    def set_gradient_checkpointing(self, checkpointing_types):
-        gradient_checkpointing = []
-        for down_block, sac_type in zip(self.down_blocks, checkpointing_types):
-            if sac_type == "coarse":
-                gradient_checkpointing.append(True)
-            elif sac_type == "fine":
-                for n, m in down_block.named_modules():
-                    if hasattr(m, "gradient_checkpointing"):
-                        m.gradient_checkpointing = True
-                        logger.debug(f"set gradient_checkpointing: {n}")
-                gradient_checkpointing.append(False)
-            else:
-                gradient_checkpointing.append(False)
-        self.gradient_checkpointing = gradient_checkpointing
-        logger.info(f"[Encoder3D] gradient_checkpointing: {checkpointing_types}")
-
-    def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
-        r"""The forward method of the `Encoder` class."""
-        sample = self.conv_in(sample, memory_state=memory_state)
-        # down
-        for down_block, sac in zip(self.down_blocks, self.gradient_checkpointing):
-            sample = gradient_checkpointing(
-                down_block,
-                sample,
-                memory_state=memory_state,
-                enabled=self.training and sac,
-            )
-
-        # middle
-        sample = self.mid_block(sample, memory_state=memory_state)
-
-        # post-process
-        sample = causal_norm_wrapper(self.conv_norm_out, sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample, memory_state=memory_state)
-
-        return sample
-
-
-class Decoder3D(nn.Module):
-    r"""
-    The `Decoder` layer of a variational autoencoder that
-    decodes its latent representation into an output sample.
-    """
-
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        block_out_channels: Tuple[int, ...] = (64,),
-        layers_per_block: int = 2,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        temporal_up_num: int = 2,
-        slicing_up_num: int = 0,
-        selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
-    ):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-        self.temporal_up_num = temporal_up_num
-
-        self.conv_in = init_causal_conv3d(
-            in_channels,
-            block_out_channels[-1],
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            inflation_mode=inflation_mode,
-        )
-
-        self.up_blocks = nn.ModuleList([])
-
-        # mid
-        self.mid_block = UNetMidBlock3D(
-            channels=block_out_channels[-1],
-            inflation_mode=inflation_mode,
-            time_receptive_field=time_receptive_field,
-        )
-
-        # up
-        reversed_block_out_channels = list(reversed(block_out_channels))
-        output_channel = reversed_block_out_channels[0]
-        for i in range(len(reversed_block_out_channels)):
-            prev_output_channel = output_channel
-            output_channel = reversed_block_out_channels[i]
-
-            is_final_block = i == len(block_out_channels) - 1
-            is_temporal_up_block = i < self.temporal_up_num
-            is_slicing_up_block = i >= len(block_out_channels) - slicing_up_num
-            # Note: Keep symmetric
-
-            up_block = UpDecoderBlock3D(
-                num_layers=self.layers_per_block + 1,
-                in_channels=prev_output_channel,
-                out_channels=output_channel,
-                add_upsample=not is_final_block,
-                temporal_up=is_temporal_up_block,
-                slicing=is_slicing_up_block,
-                inflation_mode=inflation_mode,
-                time_receptive_field=time_receptive_field,
-            )
-            self.up_blocks.append(up_block)
-
-        # out
-        self.conv_norm_out = nn.GroupNorm(
-            num_channels=block_out_channels[0], num_groups=32, eps=1e-6
-        )
-        self.conv_act = nn.SiLU()
-        self.conv_out = init_causal_conv3d(
-            block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode
-        )
-
-        assert len(selective_checkpointing) == len(self.up_blocks)
-        self.set_gradient_checkpointing(selective_checkpointing)
-
-    def set_gradient_checkpointing(self, checkpointing_types):
-        gradient_checkpointing = []
-        for up_block, sac_type in zip(self.up_blocks, checkpointing_types):
-            if sac_type == "coarse":
-                gradient_checkpointing.append(True)
-            elif sac_type == "fine":
-                for n, m in up_block.named_modules():
-                    if hasattr(m, "gradient_checkpointing"):
-                        m.gradient_checkpointing = True
-                        logger.debug(f"set gradient_checkpointing: {n}")
-                gradient_checkpointing.append(False)
-            else:
-                gradient_checkpointing.append(False)
-        self.gradient_checkpointing = gradient_checkpointing
-        logger.info(f"[Decoder3D] gradient_checkpointing: {checkpointing_types}")
-
-    def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
-        r"""The forward method of the `Decoder` class."""
-
-        sample = self.conv_in(sample, memory_state=memory_state)
-
-        # middle
-        sample = self.mid_block(sample, memory_state=memory_state)
-
-        # up
-        for up_block, sac in zip(self.up_blocks, self.gradient_checkpointing):
-            sample = gradient_checkpointing(
-                up_block,
-                sample,
-                memory_state=memory_state,
-                enabled=self.training and sac,
-            )
-
-        # post-process
-        sample = causal_norm_wrapper(self.conv_norm_out, sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample, memory_state=memory_state)
-
-        return sample
-
-
-class VideoAutoencoderKL(nn.Module):
-    def __init__(
-        self,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        block_out_channels: Tuple[int] = (64,),
-        layers_per_block: int = 1,
-        latent_channels: int = 4,
-        use_quant_conv: bool = True,
-        use_post_quant_conv: bool = True,
-        enc_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
-        dec_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
-        temporal_scale_num: int = 0,
-        slicing_up_num: int = 0,
-        inflation_mode: _inflation_mode_t = "tail",
-        time_receptive_field: _receptive_field_t = "half",
-        slicing_sample_min_size: int = None,
-        spatial_downsample_factor: int = 16,
-        temporal_downsample_factor: int = 8,
-        freeze_encoder: bool = False,
-    ):
-        super().__init__()
-        self.spatial_downsample_factor = spatial_downsample_factor
-        self.temporal_downsample_factor = temporal_downsample_factor
-        self.freeze_encoder = freeze_encoder
-        if slicing_sample_min_size is None:
-            slicing_sample_min_size = temporal_downsample_factor
-        self.slicing_sample_min_size = slicing_sample_min_size
-        self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
-
-        # pass init params to Encoder
-        self.encoder = Encoder3D(
-            in_channels=in_channels,
-            out_channels=latent_channels,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            double_z=True,
-            temporal_down_num=temporal_scale_num,
-            selective_checkpointing=enc_selective_checkpointing,
-            inflation_mode=inflation_mode,
-            time_receptive_field=time_receptive_field,
-        )
-
-        # pass init params to Decoder
-        self.decoder = Decoder3D(
-            in_channels=latent_channels,
-            out_channels=out_channels,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            # [Override] add temporal_up_num parameter
-            temporal_up_num=temporal_scale_num,
-            slicing_up_num=slicing_up_num,
-            selective_checkpointing=dec_selective_checkpointing,
-            inflation_mode=inflation_mode,
-            time_receptive_field=time_receptive_field,
-        )
-
-        self.quant_conv = (
-            init_causal_conv3d(
-                in_channels=2 * latent_channels,
-                out_channels=2 * latent_channels,
-                kernel_size=1,
-                inflation_mode=inflation_mode,
-            )
-            if use_quant_conv
-            else None
-        )
-        self.post_quant_conv = (
-            init_causal_conv3d(
-                in_channels=latent_channels,
-                out_channels=latent_channels,
-                kernel_size=1,
-                inflation_mode=inflation_mode,
-            )
-            if use_post_quant_conv
-            else None
-        )
-
-        self.use_slicing = False
-
-    def enable_slicing(self):
-        self.use_slicing = True
-
-    def disable_slicing(self):
-        self.use_slicing = False
-
-    def encode(self, x: torch.FloatTensor) -> CausalEncoderOutput:
-        if x.ndim == 4:
-            x = x.unsqueeze(2)
-        h = self.slicing_encode(x)
-        p = DiagonalGaussianDistribution(h)
-        z = p.sample()
-        return CausalEncoderOutput(z, p)
-
-    def decode(self, z: torch.FloatTensor) -> CausalDecoderOutput:
-        if z.ndim == 4:
-            z = z.unsqueeze(2)
-        x = self.slicing_decode(z)
-        return CausalDecoderOutput(x)
-
-    def _encode(self, x: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
-        x = causal_conv_slice_inputs(x, self.slicing_sample_min_size, memory_state=memory_state)
-        h = self.encoder(x, memory_state=memory_state)
-        h = self.quant_conv(h, memory_state=memory_state) if self.quant_conv is not None else h
-        h = causal_conv_gather_outputs(h)
-        return h
-
-    def _decode(self, z: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
-        z = causal_conv_slice_inputs(z, self.slicing_latent_min_size, memory_state=memory_state)
-        z = (
-            self.post_quant_conv(z, memory_state=memory_state)
-            if self.post_quant_conv is not None
-            else z
-        )
-        x = self.decoder(z, memory_state=memory_state)
-        x = causal_conv_gather_outputs(x)
-        return x
-
-    def slicing_encode(self, x: torch.Tensor) -> torch.Tensor:
-        sp_size = 1
-        if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size * sp_size:
-            x_slices = x[:, :, 1:].split(split_size=self.slicing_sample_min_size * sp_size, dim=2)
-            encoded_slices = [
-                self._encode(
-                    torch.cat((x[:, :, :1], x_slices[0]), dim=2),
-                    memory_state=MemoryState.INITIALIZING,
-                )
-            ]
-            for x_idx in range(1, len(x_slices)):
-                encoded_slices.append(
-                    self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE)
-                )
-            return torch.cat(encoded_slices, dim=2)
-        else:
-            return self._encode(x, memory_state=MemoryState.DISABLED)
-
-    def slicing_decode(self, z: torch.Tensor) -> torch.Tensor:
-        sp_size = 1
-        if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size * sp_size:
-            z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size * sp_size, dim=2)
-            decoded_slices = [
-                self._decode(
-                    torch.cat((z[:, :, :1], z_slices[0]), dim=2),
-                    memory_state=MemoryState.INITIALIZING,
-                )
-            ]
-            for z_idx in range(1, len(z_slices)):
-                decoded_slices.append(
-                    self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE)
-                )
-            return torch.cat(decoded_slices, dim=2)
-        else:
-            return self._decode(z, memory_state=MemoryState.DISABLED)
-
-    def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:
-        with torch.no_grad() if self.freeze_encoder else nullcontext():
-            z, p = self.encode(x)
-        x = self.decode(z).sample
-        return CausalAutoencoderOutput(x, z, p)
-
-    def preprocess(self, x: torch.Tensor):
-        # x should in [B, C, T, H, W], [B, C, H, W]
-        assert x.ndim == 4 or x.size(2) % self.temporal_downsample_factor == 1
-        return x
-
-    def postprocess(self, x: torch.Tensor):
-        # x should in [B, C, T, H, W], [B, C, H, W]
-        return x
-
-    def set_causal_slicing(
-        self,
-        *,
-        split_size: Optional[int],
-        memory_device: _memory_device_t,
-    ):
-        assert (
-            split_size is None or memory_device is not None
-        ), "if split_size is set, memory_device must not be None."
-        if split_size is not None:
-            self.enable_slicing()
-            self.slicing_sample_min_size = split_size
-            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
-        else:
-            self.disable_slicing()
-        for module in self.modules():
-            if isinstance(module, InflatedCausalConv3d):
-                module.set_memory_device(memory_device)
-
-    def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]):
-        set_norm_limit(norm_max_mem)
-        for m in self.modules():
-            if isinstance(m, InflatedCausalConv3d):
-                m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf"))
-
-
-class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
-    def __init__(
-        self, *args, spatial_downsample_factor: int, temporal_downsample_factor: int, **kwargs
-    ):
-        self.spatial_downsample_factor = spatial_downsample_factor
-        self.temporal_downsample_factor = temporal_downsample_factor
-        super().__init__(*args, **kwargs)
-
-    def forward(self, x) -> CausalAutoencoderOutput:
-        z, _, p = self.encode(x)
-        x, _ = self.decode(z)
-        return CausalAutoencoderOutput(x, z, None, p)
-
-    def encode(self, x) -> CausalEncoderOutput:
-        if x.ndim == 4:
-            x = x.unsqueeze(2)
-        p = super().encode(x).latent_dist
-        z = p.sample().squeeze(2)
-        return CausalEncoderOutput(z, None, p)
-
-    def decode(self, z) -> CausalDecoderOutput:
-        if z.ndim == 4:
-            z = z.unsqueeze(2)
-        x = super().decode(z).sample.squeeze(2)
-        return CausalDecoderOutput(x, None)
-
-    def preprocess(self, x):
-        # x should in [B, C, T, H, W], [B, C, H, W]
-        assert x.ndim == 4 or x.size(2) % 4 == 1
-        return x
-
-    def postprocess(self, x):
-        # x should in [B, C, T, H, W], [B, C, H, W]
-        return x
-
-    def set_causal_slicing(
-        self,
-        *,
-        split_size: Optional[int],
-        memory_device: Optional[Literal["cpu", "same"]],
-    ):
-        assert (
-            split_size is None or memory_device is not None
-        ), "if split_size is set, memory_device must not be None."
-        if split_size is not None:
-            self.enable_slicing()
-        else:
-            self.disable_slicing()
-        self.slicing_sample_min_size = split_size
-        if split_size is not None:
-            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
-        for module in self.modules():
-            if isinstance(module, InflatedCausalConv3d):
-                module.set_memory_device(memory_device)

modules/seedvr/src/optimization/memory_manager.py

@@ -84,27 +84,12 @@ def clear_rope_cache(runner) -> None:
         runner: The model runner containing the cache
     """
     if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'):
-        # Count entries before cleanup
-        len(runner.cache.cache)
-
-        # Free all tensors from cache
-        for _key, value in runner.cache.cache.items():
-            if isinstance(value, (tuple, list)):
-                for item in value:
-                    if hasattr(item, 'cpu'):
-                        item.cpu()
-                        del item
-            elif hasattr(value, 'cpu'):
-                value.cpu()
-                del value
-
-        # Clear the cache
-        runner.cache.cache.clear()
+        runner.cache.clear()
 
     if hasattr(runner, 'dit'):
-        cleared_lru_count = 0
         for module in runner.dit.modules():
             if isinstance(module, RotaryEmbeddingBase):
                 if hasattr(module.get_axial_freqs, 'cache_clear'):
                     module.get_axial_freqs.cache_clear()
-                    cleared_lru_count += 1
+            if hasattr(module, 'cache') and hasattr(module.cache, 'clear'):
+                module.cache.clear()

modules/seedvr/src/optimization/performance.py

@@ -54,7 +54,7 @@ def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.T
         batch_3d = batch_3d.permute(0, 2, 1, 3, 4)  # [batch, 1, c, h, w]
 
         for i, idx in enumerate(indices_3d):
-            samples[idx] = batch_3d[i]  # [1, c, h, w]
+            samples[idx] = batch_3d[i].contiguous()  # [1, c, h, w]
 
     # 🚀 BATCH PROCESSING for 4D videos (c t h w -> t c h w)
     if videos_4d:
@@ -67,13 +67,12 @@ def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.T
             batch_4d = batch_4d.permute(0, 2, 1, 3, 4)  # [batch, t, c, h, w]
 
             for i, idx in enumerate(indices_4d):
-                samples[idx] = batch_4d[i]  # [t, c, h, w]
+                samples[idx] = batch_4d[i].contiguous()  # [t, c, h, w]
         else:
             # 🔄 FALLBACK: Different shapes, optimized individual processing
             for i, idx in enumerate(indices_4d):
                 # Use permute instead of rearrange (faster)
-                samples[idx] = videos_4d[i].permute(1, 0, 2, 3)  # c t h w -> t c h w
-
+                samples[idx] = videos_4d[i].permute(1, 0, 2, 3).contiguous()  # c t h w -> t c h w
     return samples
 
 

modules/seedvr/src/utils/color_fix.py

@@ -2,8 +2,9 @@ import torch
 from PIL import Image
 from torch import Tensor
 from torch.nn import functional as F
-from modules.seedvr.src.common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation
 from torchvision.transforms import ToTensor, ToPILImage
+from modules.seedvr.src.common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation
+
 
 def adain_color_fix(target: Image.Image, source: Image.Image):
     # Convert images to tensors
@@ -118,6 +119,10 @@ def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
                 align_corners=False
             )
 
+    # align devices so reconstruction does not mix CPU and GPU tensors
+    if style_feat.device != content_feat.device:
+        style_feat = style_feat.to(content_feat.device)
+
     # calculate the wavelet decomposition of the content feature
     content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
     del content_low_freq

modules/seedvr/test.py

@@ -36,6 +36,8 @@ def upscale_image(model_name:str, image_path:str):
         runner=runner,
         images=image_tensor,
         cfg_scale=cfg,
+        cfg_rescale=0.0,
+        steps=1,
         seed=seed,
         res_w=resolution,
         batch_size=1,

modules/shared_items.py

@@ -180,4 +180,4 @@ def get_repo(model):
 
 
 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"]
+sdnq_matmul_modes = ["disabled", "enabled", "int8", "uint8", "float8_e4m3fn", "float16"]

modules/shared_state.py

@@ -38,7 +38,7 @@ class State:
     current_sigma = None
     current_sigma_next = None
     current_image = None
-    current_image_sampling_step = 0
+    current_image_sampling_step = -1
     id_live_preview = 0
     textinfo = None
     prediction_type = "epsilon"
@@ -92,7 +92,7 @@ class State:
         self.do_set_current_image()
         self.job_no += 1
         # self.sampling_step = 0
-        self.current_image_sampling_step = 0
+        self.current_image_sampling_step = -1
         if debug_output:
             log.trace(f'State next: {self}')
         modules.devices.torch_gc()
@@ -202,7 +202,7 @@ class State:
         self.job_history += 1
         self.total_jobs += 1
         self.current_image = None
-        self.current_image_sampling_step = 0
+        self.current_image_sampling_step = -1
         self.current_latent = None
         self.current_noise_pred = None
         self.current_sigma = None
@@ -271,7 +271,7 @@ class State:
 
     def do_set_current_image(self):
         from modules import shared, images, sd_samplers_common
-        if self.disable_preview or (self.preview_job == self.job_no):
+        if self.disable_preview or (self.preview_job == self.job_no) or (self.current_image_sampling_step == self.sampling_step):
             return False
 
         if (shared.opts.show_progress_type == "None") and (shared.history.last_image is not None):

modules/styles.py

@@ -247,6 +247,9 @@ def apply_styles_to_extra(p, style: Style):
     extra.update(infotext.parse(style_extra))
     extra.pop('Prompt', None)
     extra.pop('Negative prompt', None)
+    if debug_enabled:
+        log.trace(f'Apply style extra: {extra}')
+
     params = []
     settings = []
     skipped = []

modules/timer.py

@@ -10,11 +10,11 @@ except Exception:
 
 
 class Timer:
-    def __init__(self):
+    def __init__(self, profile=False):
         self.start = time.time()
         self.records = {}
         self.total = 0
-        self.profile = False
+        self.profile = profile
 
     def elapsed(self, reset=True):
         end = time.time()
@@ -28,6 +28,13 @@ class Timer:
             self.records[name] = 0
         self.records[name] += t
 
+    def rm(self, name):
+        if name in self.records:
+            del self.records[name]
+
+    def get(self, name):
+        return self.records.get(name, 0)
+
     def ts(self, name, t):
         elapsed = time.time() - t
         self.add(name, elapsed)
@@ -60,9 +67,12 @@ class Timer:
     def get_total(self):
         return sum(self.records.values())
 
-    def dct(self, min_time=default_min_time):
+    def dct(self, min_time=default_min_time, no_total=False):
         self.total = sum(self.records.values())
-        self.records['total'] = self.total
+        if no_total:
+            self.records.pop('total', None)
+        else:
+            self.records['total'] = self.total
         if self.profile:
             res = {k: round(v, 4) for k, v in self.records.items()}
         else:
@@ -71,10 +81,13 @@ class Timer:
         return res
 
     def reset(self):
-        self.__init__()
+        self.records.clear()
+        self.__init__(self.profile)
 
 startup = Timer()
 process = Timer()
 launch = Timer()
 init = Timer()
 load = Timer()
+dynamo = Timer()
+autotune = Timer(profile=True)

modules/timer_sdnq.py

@@ -0,0 +1,13 @@
+def update_sdnq_attention_timers():
+    from modules.timer import autotune
+    autotune.reset()
+    from modules.sdnq.kernels import triton_atten, triton_mm, triton_scaled_mm
+    if getattr(triton_atten.sdnq_attn_kernel, 'bench_time', None) is not None:
+        autotune.add('sdnq_attn_kernel', getattr(triton_atten.sdnq_attn_kernel, 'bench_time', 0))
+        triton_atten.sdnq_attn_kernel.bench_time = 0
+    if getattr(triton_mm.sdnq_triton_mm, 'bench_time', None) is not None:
+        autotune.add('sdnq_triton_mm', getattr(triton_mm.sdnq_triton_mm, 'bench_time', 0))
+        triton_mm.sdnq_triton_mm.bench_time = 0
+    if getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', None) is not None:
+        autotune.add('sdnq_scaled_mm', getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', 0))
+        triton_scaled_mm.sdnq_scaled_mm.bench_time = 0

modules/transformer_cache.py

@@ -19,15 +19,15 @@ def set_cache(faster_cache=None, pyramid_attention_broadcast=None):
         return
     try:
         if faster_cache: # https://github.com/huggingface/diffusers/pull/10163
-            distilled = shared.opts.fc_guidance_distilled or shared.sd_model_type == 'f1'
+            distilled = shared.opts.fc_guidance_distilled
             config = diffusers.FasterCacheConfig(
                 spatial_attention_block_skip_range=shared.opts.fc_spacial_skip_range,
                 spatial_attention_timestep_skip_range=(int(shared.opts.fc_spacial_skip_start), int(shared.opts.fc_spacial_skip_end)),
                 unconditional_batch_skip_range=shared.opts.fc_uncond_skip_range,
                 unconditional_batch_timestep_skip_range=(int(shared.opts.fc_uncond_skip_start), int(shared.opts.fc_uncond_skip_end)),
                 attention_weight_callback=lambda _: shared.opts.fc_attention_weight,
-                tensor_format=shared.opts.fc_tensor_format, # TODO fc: autodetect tensor format based on model
-                is_guidance_distilled=distilled, # TODO fc: autodetect distilled based on model
+                tensor_format=shared.opts.fc_tensor_format,
+                is_guidance_distilled=distilled,
                 current_timestep_callback=lambda: shared.sd_model.current_timestep,
             )
             shared.sd_model.transformer.disable_cache()

modules/ui_common.py

@@ -62,14 +62,18 @@ def infotext_to_html(text):
     res.pop('Negative template', None)
 
     runtime = {}
-    runtime['App'] = res.get('App', '')
-    res.pop('App', None)
-    runtime['Version'] = res.get('Version', '')
-    res.pop('Version', None)
-    runtime['Pipeline'] = res.get('Pipeline', '')
-    res.pop('Pipeline', None)
-    runtime['Operations'] = res.get('Operations', '')
-    res.pop('Operations', None)
+    if 'App' in res:
+        runtime['App'] = res.get('App', '')
+        res.pop('App', None)
+    if 'Version' in res:
+        runtime['Version'] = res.get('Version', '')
+        res.pop('Version', None)
+    if 'Pipeline' in res:
+        runtime['Pipeline'] = res.get('Pipeline', '')
+        res.pop('Pipeline', None)
+    if 'Operations' in res:
+        runtime['Operations'] = res.get('Operations', '')
+        res.pop('Operations', None)
 
     params = [f'{k}: {v}' for k, v in res.items() if v is not None and not k.endswith('-1') and not k.endswith('-2')]
     params = '| '.join(params) if len(params) > 0 else ''

modules/ui_definitions.py

@@ -174,9 +174,9 @@ def create_settings(cmd_opts):
         "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"]}),
         "sdnq_quantize_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"]}),
         "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes}),
-        "sdnq_quantize_matmul_mode": OptionInfo("auto", "Quantized MatMul type", gr.Dropdown, {"choices": sdnq_matmul_modes}),
+        "sdnq_quantize_matmul_mode": OptionInfo("disabled", "Quantized MatMul type", gr.Dropdown, {"choices": sdnq_matmul_modes}),
         "sdnq_quantize_weights_mode_te": OptionInfo("Same as model", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_quant_modes}),
-        "sdnq_quantize_matmul_mode_te": OptionInfo("Same as model", "Quantized MatMul type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_matmul_modes}),
+        "sdnq_quantize_matmul_mode_te": OptionInfo("disabled", "Quantized MatMul type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_matmul_modes}),
         "sdnq_modules_to_not_convert": OptionInfo("", "Modules to not convert"),
         "sdnq_modules_dtype_dict": OptionInfo("{}", "Modules dtype dict"),
         "sdnq_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}),
@@ -190,7 +190,6 @@ def create_settings(cmd_opts):
         "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox),
         "sdnq_quantize_embedding_layers": OptionInfo(False, "Quantize embedding layers", gr.Checkbox),
         "sdnq_dequantize_compile": OptionInfo(devices.has_triton(early=True), "Dequantize using torch.compile", gr.Checkbox),
-        "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox),
         "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use quantized MatMul with conv", gr.Checkbox),
         "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize using GPU", gr.Checkbox),
         "sdnq_dequantize_fp32": OptionInfo(True, "Dequantize using full precision", gr.Checkbox),
@@ -260,9 +259,8 @@ def create_settings(cmd_opts):
         "sdnq_attention_sep": OptionInfo("<h2>SDNQ Attention</h2>", "", gr.HTML),
         "sdnq_attention_smooth_k": OptionInfo(False, "SDNQ Attention use Smooth K", gr.Checkbox),
         "sdnq_attention_use_hadamard": OptionInfo(False, "SDNQ Attention use Hadamard", gr.Checkbox),
-        "sdnq_attention_use_quantized_matmul": OptionInfo(True, "SDNQ Attention use Quantized MatMul", gr.Checkbox),
-        "sdnq_attention_matmul_type": OptionInfo("auto", "SDNQ Attention MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
-        "sdnq_attention_pv_matmul_type": OptionInfo("auto", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
+        "sdnq_attention_matmul_type": OptionInfo("enabled", "SDNQ Attention MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
+        "sdnq_attention_pv_matmul_type": OptionInfo("disabled", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
         "sdnq_attention_hadamard_group_size": OptionInfo(256, "SDNQ Attention Hadamard Group Size", gr.Slider, {"minimum": 4, "maximum": 1024, "step": 1}),
 
         "hf_attention_sep": OptionInfo("<h2>Attention Dispatcher</h2>", "", gr.HTML),
@@ -633,9 +631,6 @@ def create_settings(cmd_opts):
         "detailer_unload": OptionInfo(False, "Move detailer model to CPU when complete"),
         "detailer_augment": OptionInfo(False, "Detailer use model augment"),
 
-        "postprocessing_sep_seedvr": OptionInfo("<h2>SeedVR</h2>", "", gr.HTML),
-        "seedvr_cfg_scale": OptionInfo(3.5, "SeedVR CFG Scale", gr.Slider, {"minimum": 1, "maximum": 15, "step": 1}),
-
         "postprocessing_sep_upscalers": OptionInfo("<h2>Upscaling</h2>", "", gr.HTML),
         "upscaler_unload": OptionInfo(False, "Unload upscaler after processing"),
         "upscaler_latent_steps": OptionInfo(20, "Upscaler latent steps", gr.Slider, {"minimum": 4, "maximum": 100, "step": 1}),

modules/ui_extra_networks_styles.py

@@ -77,8 +77,8 @@ class ExtraNetworksPageStyles(ui_extra_networks.ExtraNetworksPage):
             name = getattr(style, 'name', '')
             if name == '':
                 return item
-            txt = f'Prompt: {getattr(style, "prompt", "")}'
-            if len(getattr(style, 'negative_prompt', '')) > 0:
+            txt = f'Prompt: {getattr(style, "prompt", "") or ""}'
+            if len(getattr(style, 'negative_prompt', '') or '') > 0:
                 txt += f'\nNegative: {style.negative_prompt}'
             item = {
                 "type": 'Style',

modules/ui_postprocessing.py

@@ -1,5 +1,6 @@
+import os
 import gradio as gr
-from modules import scripts_manager, shared, ui_common, postprocessing, call_queue, generation_parameters_copypaste
+from modules import scripts_manager, shared, progress, ui_common, postprocessing, call_queue, generation_parameters_copypaste
 from modules.logger import log
 
 
@@ -12,10 +13,32 @@ def submit_info(image):
     return infotext_to_html(geninfo), info, geninfo
 
 
-def submit_process(tab_index, extras_image, image_batch, extras_batch_input_dir, extras_batch_output_dir, show_extras_results, save_output, *script_inputs):
+def submit_video(video):
+    if not video or not isinstance(video, str) or not os.path.isfile(video):
+        return '', '', ''
+    from modules.video import get_video_info, get_video_metadata
+    info = get_video_info(video)
+    metadata = get_video_metadata(video)
+    if metadata:
+        info['metadata'] = metadata
+    text = ''
+    html = ''
+    html = [f'<b>{k}</b>: {v}' for k, v in info.items()]
+    html = '<br>'.join(html)
+    text = [f'{k}: {v}' for k, v in info.items()]
+    text = ', '.join(text)
+    return html, '', text
+
+
+def submit_process(job_id, tab_index, extras_image, image_batch, extras_batch_input_dir, extras_batch_output_dir, extras_video, show_extras_results, save_output, *script_inputs):
+    progress.start_task(job_id)
     from modules.ui_common import infotext_to_html
-    result_images, geninfo, _js_info = postprocessing.run_postprocessing(tab_index, extras_image, image_batch, extras_batch_input_dir, extras_batch_output_dir, show_extras_results, *script_inputs, save_output=save_output)
-    return result_images, geninfo, infotext_to_html(geninfo)
+    result_images, result_video, geninfo, _js_info = postprocessing.run_postprocessing(tab_index, extras_image, image_batch, extras_batch_input_dir, extras_batch_output_dir, extras_video, show_extras_results, *script_inputs, save_output=save_output)
+    progress.finish_task(job_id)
+
+    gr_result_image = gr.update(value=result_images, visible=tab_index != 3)
+    gr_result_video = gr.update(value=result_video, visible=tab_index == 3)
+    return gr_result_image, gr_result_video, geninfo, infotext_to_html(geninfo)
 
 
 def create_ui():
@@ -25,31 +48,39 @@ def create_ui():
         with gr.Column(variant='compact'):
             with gr.Tabs(elem_id="mode_extras"):
                 with gr.Tab('Process Image', id="single_image", elem_id="extras_single_tab") as tab_single:
-                    with gr.Row():
-                        extras_image = gr.Image(label="Source", interactive=True, type="pil", elem_id="extras_image")
+                    extras_image = gr.Image(label="Source", interactive=True, type="pil", elem_id="extras_image")
                 with gr.Tab('Process Batch', id="batch_process", elem_id="extras_batch_process_tab") as tab_batch:
                     image_batch = gr.Files(label="Batch process", interactive=True, elem_id="extras_image_batch")
                 with gr.Tab('Process Folder', id="batch_from_directory", elem_id="extras_batch_directory_tab") as tab_batch_dir:
                     extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, placeholder="A directory on the same machine where the server is running.", elem_id="extras_batch_input_dir")
                     extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, placeholder="Leave blank to save images to the default path.", elem_id="extras_batch_output_dir")
                     show_extras_results = gr.Checkbox(label='Show result images', value=True, elem_id="extras_show_extras_results")
+                with gr.Tab('Process Video', id="process_video", elem_id="extras_process_video_tab") as tab_process_video:
+                    extras_video = gr.Video(label="Input Video", show_label=False, height=512, interactive=True, elem_id="extras_video")
             with gr.Row():
                 save_output = gr.Checkbox(label='Save output', value=True, elem_id="extras_save_output")
 
             script_inputs = scripts_manager.scripts_postproc.setup_ui()
-        with gr.Column():
+        with gr.Column(elem_id="extras_output_column"):
             id_part = 'extras'
             with gr.Row(elem_id=f"{id_part}_generate_box", elem_classes="generate-box"):
-                submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')
+                submit = gr.Button('Process', elem_id=f"{id_part}_generate", variant='primary')
                 interrupt = gr.Button('Stop', elem_id=f"{id_part}_interrupt", variant='secondary')
                 interrupt.click(fn=shared.state.interrupt, inputs=[], outputs=[])
                 skip = gr.Button('Skip', elem_id=f"{id_part}_skip", variant='secondary')
                 skip.click(fn=shared.state.skip, inputs=[], outputs=[])
                 pause = gr.Button('Pause', elem_id=f"{id_part}_pause")
                 pause.click(fn=shared.state.pause, _js='checkPaused', inputs=[], outputs=[])
-            result_images, generation_info, _html_info, html_info_formatted, _html_log = ui_common.create_output_panel("extras")
+
+            with gr.Tabs(elem_id="extras_output_tabs"):
+                with gr.Tab('Image', id="process_output_image", elem_id="extras_output_image_tab"):
+                    result_images, generation_info, _html_info, html_info_formatted, _html_log = ui_common.create_output_panel("extras")
+                with gr.Tab('Video', id="process_output_video", elem_id="extras_output_video_tab"):
+                    result_video = gr.Video(label="Video", show_label=False, interactive=False, elem_id="extras_output_video", visible=False)
+
             gr.HTML('File metadata')
             exif_info = gr.HTML(elem_id="pnginfo_html_info")
+
             with gr.Row(elem_id='copy_buttons_process'):
                 copy_process_buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "caption"])
 
@@ -57,25 +88,33 @@ def create_ui():
             generation_parameters_copypaste.register_paste_params_button(generation_parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=generation_info, source_image_component=extras_image))
         generation_parameters_copypaste.add_paste_fields("extras", extras_image, None)
 
+    job_id = gr.Textbox(value='none', visible=False)
     tab_single.select(fn=lambda: 0, inputs=[], outputs=[tab_index])
     tab_batch.select(fn=lambda: 1, inputs=[], outputs=[tab_index])
     tab_batch_dir.select(fn=lambda: 2, inputs=[], outputs=[tab_index])
+    tab_process_video.select(fn=lambda: 3, inputs=[], outputs=[tab_index])
+
     extras_image.change(fn=submit_info, inputs=[extras_image], outputs=[html_info_formatted, exif_info, generation_info])
+    extras_video.change(fn=submit_video, inputs=[extras_video], outputs=[html_info_formatted, exif_info, generation_info])
+
     submit.click(
         _js="submit_postprocessing",
-        fn=call_queue.wrap_gradio_gpu_call(submit_process, extra_outputs=[None, ''], name='Postprocess'),
+        fn=call_queue.wrap_gradio_gpu_call(submit_process, extra_outputs=[None, None, ''], name='Postprocess'),
         inputs=[
+            job_id,
             tab_index,
             extras_image,
             image_batch,
             extras_batch_input_dir,
             extras_batch_output_dir,
+            extras_video,
             show_extras_results,
             save_output,
             *script_inputs,
         ],
         outputs=[
             result_images,
+            result_video,
             generation_info,
             html_info_formatted,
         ]

modules/video.py

@@ -105,7 +105,10 @@ def get_video_params(filepath: str, capture: bool = False):
         raise RuntimeError(msg)
     frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
     fps = round(video.get(cv2.CAP_PROP_FPS), 2)
-    duration = round(float(frames) / fps, 2)
+    if fps > 0:
+        duration = round(float(frames) / fps, 2)
+    else:
+        duration = 0
     w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
     codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC))
     frame = None
@@ -115,3 +118,92 @@ def get_video_params(filepath: str, capture: bool = False):
         frame = Image.fromarray(frame)
     video.release()
     return frames, fps, duration, w, h, codec, frame
+
+
+def get_video_info(filepath: str):
+    import cv2
+    from modules.control.util import decode_fourcc
+
+    try:
+        info = {
+            'file': os.path.basename(filepath),
+            'size': os.path.getsize(filepath),
+            'container': os.path.splitext(filepath)[1].lower().lstrip('.'),
+        }
+    except Exception as e:
+        log.error(f'Video probe failed: path="{filepath}" {e}')
+        return {}
+
+    def get_prop(video, name: str, label: str | None, cast: type = float):
+        prop_id = getattr(cv2, name, None)
+        if prop_id is None:
+            return None
+        try:
+            value = video.get(prop_id)
+            if value is None:
+                return None
+            value = cast(value)
+            if value == 0:
+                return None
+            if label is not None:
+                info[label] = value
+            return value
+        except Exception:
+            return None
+
+    video = cv2.VideoCapture(filepath)
+    try:
+        if not video.isOpened():
+            msg = f'Video open failed: path="{filepath}"'
+            info['error'] = msg
+            log.error(msg)
+            return info
+        try:
+            backend = video.getBackendName()
+            if backend:
+                info['backend'] = backend
+        except Exception:
+            pass
+
+        frames = get_prop(video, 'CAP_PROP_FRAME_COUNT', 'frames', int) or 0
+        fps = get_prop(video, 'CAP_PROP_FPS', 'fps', float) or 0
+        _width = get_prop(video, 'CAP_PROP_FRAME_WIDTH', 'width', int) or 0
+        _height = get_prop(video, 'CAP_PROP_FRAME_HEIGHT', 'height', int) or 0
+        _bitrate = get_prop(video, 'CAP_PROP_BITRATE', 'bitrate_kbps', float)
+        _orientation = get_prop(video, 'CAP_PROP_ORIENTATION_META', 'rotation', float)
+
+        if frames > 0 and fps > 0:
+            info['duration'] = round(float(frames) / fps, 3)
+
+        fourcc = get_prop(video, 'CAP_PROP_FOURCC', None, int) or 0
+        info['codec'] = decode_fourcc(fourcc)
+
+        pix = get_prop(video, 'CAP_PROP_CODEC_PIXEL_FORMAT', None, int) or 0
+        info['format'] = decode_fourcc(pix)
+
+        sar_num = get_prop(video, 'CAP_PROP_SAR_NUM', None, int) or 1
+        sar_den = get_prop(video, 'CAP_PROP_SAR_DEN', None, int) or 1
+        if (sar_num > 0 and sar_den > 0) and (sar_num != 1 or sar_den != 1):
+            info['sar'] = f'{sar_num}/{sar_den}'
+
+        audio_streams = get_prop(video, 'CAP_PROP_AUDIO_TOTAL_STREAMS', 'audio_streams', int) or 0
+        audio_channels = get_prop(video, 'CAP_PROP_AUDIO_TOTAL_CHANNELS', 'audio_channels', int) or 0
+        audio_sample_rate = get_prop(video, 'CAP_PROP_AUDIO_SAMPLES_PER_SECOND', 'audio_sample_rate', int) or 0
+        if audio_streams > 0 or audio_channels > 0 or audio_sample_rate > 0:
+            info['audio'] = f'{audio_streams}x{audio_channels}x{audio_sample_rate}'
+
+        return info
+    except Exception as e:
+        log.error(f'Video probe failed: path="{filepath}" {e}')
+        return info
+    finally:
+        video.release()
+
+
+def get_video_metadata(video):
+    if not video or not isinstance(video, str) or not os.path.isfile(video):
+        return {}
+    from modules.video_models.video_utils import check_av
+    av = check_av()
+    with av.open(video, mode="r") as container:
+        return dict(container.metadata)

modules/video_models/video_load.py

@@ -155,9 +155,7 @@ def load_model(selected: models_def.Model):
             shared.sd_model = load_custom(selected.repo)
         else:
             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,

modules/video_models/video_save.py

@@ -256,6 +256,7 @@ def save_video(
     stream=None,  # async progress reporting stream
     metadata: dict | None = None,  # metadata for video
     pbar=None,  # progress bar for video
+    reclamp: bool = True,  # reclamp pixels to [-1, 1] range
 ):
     if metadata is None:
         metadata = {}
@@ -288,6 +289,8 @@ def save_video(
         log.error(f'Video: type={type(pixels)} not a tensor')
         return 0, output_video, None
     t_save = time.time()
+    if pixels.ndim == 4:
+        pixels = pixels.unsqueeze(0)
     n, _c, t, h, w = pixels.shape
     size = pixels.element_size() * pixels.numel()
     log.debug(f'Video: video={mp4_video} export={mp4_frames} safetensors={mp4_sf} interpolate={mp4_interpolate}')
@@ -304,8 +307,10 @@ def save_video(
             pixels = pixels.permute(1, 2, 0, 3, 4)
             pixels = pixels * 2.0 - 1.0
 
-        n, _c, t, h, w = pixels.shape
-        x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
+        if reclamp:
+            x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
+        else:
+            x = pixels.float() * 255.0
         x = x.detach().cpu().to(torch.uint8)
         x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=n)
         x = x.contiguous()

pipelines/generic_shared.py

@@ -9,6 +9,7 @@ shared_te_map = {
         'cls': transformers.T5EncoderModel,
         'identifier': 'sdnq-uint4',
         'target_repo': 'Disty0/FLUX.1-dev-SDNQ-uint4-svd-r32',
+        'target_subfolder': 'text_encoder_2',
     },
     'T5-XXL Base': { # template
         'cls': transformers.T5EncoderModel, # desired model class, used as primary matching criteria
@@ -33,13 +34,7 @@ shared_te_map = {
 
     'Qwen-2.5 SDNQ-4Bit': {
         'cls': transformers.Qwen2_5_VLForConditionalGeneration,
-        'identifier': 'sdnq-4bit',
-        'target_repo': 'Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32',
-        'target_subfolder': 'text_encoder',
-    },
-    'Qwen-2.5 SDNQ-UInt4': {
-        'cls': transformers.Qwen2_5_VLForConditionalGeneration,
-        'identifier': 'sdnq-uint4',
+        'identifier': ['sdnq-4bit', 'sdnq-uint4'],
         'target_repo': 'Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32',
         'target_subfolder': 'text_encoder',
     },
@@ -51,13 +46,7 @@ shared_te_map = {
 
     'Qwen-3 9B SDNQ-4bit': {
         'cls': transformers.Qwen3ForCausalLM,
-        'identifier': '9b-sdnq-4bit',
-        'target_repo': 'Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32',
-        'target_subfolder': 'text_encoder',
-    },
-    'Qwen-3 9B SDNQ-UInt4': {
-        'cls': transformers.Qwen3ForCausalLM,
-        'identifier': '9b-sdnq-uint4',
+        'identifier': ['9b-sdnq-4bit', '9b-sdnq-uint4', '9b-sdnq-hadamard-uint4', '9b-kv-merge-sdnq-hadamard-uint4'],
         'target_repo': 'Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32',
         'target_subfolder': 'text_encoder',
     },
@@ -68,15 +57,9 @@ shared_te_map = {
         'target_subfolder': 'text_encoder',
     },
 
-    'Qwen-3 4B SDNQ-4Bit': { # match after 9b
+    'Qwen-3 4B SDNQ-4Bit': {
         'cls': transformers.Qwen3ForCausalLM,
-        'identifier': 'sdnq-4bit',
-        'target_repo': 'Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32',
-        'target_subfolder': 'text_encoder',
-    },
-    'Qwen-3 4B SDNQ-UInt4': {
-        'cls': transformers.Qwen3ForCausalLM,
-        'identifier': 'sdnq-uint4',
+        'identifier': ['sdnq-4bit', 'sdnq-uint4'],
         'target_repo': 'Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32',
         'target_subfolder': 'text_encoder',
     },
@@ -103,6 +86,7 @@ shared_te_map = {
         'identifier': 'krea',
         'target_repo': 'Qwen/Qwen3-VL-4B-Instruct',
     },
+
     'Qwen3-VL 8B SDNQ-UInt4': {
         'cls': transformers.Qwen3VLModel,
         'identifier': 'uint4',
@@ -114,6 +98,18 @@ shared_te_map = {
         'target_repo': 'Qwen/Qwen3-VL-8B-Instruct',
     },
 
+    'Qwen3-VL 2B Conditional': {
+        'cls': transformers.Qwen3VLForConditionalGeneration,
+        'target_repo': 'SeFi-Image/SeFi-Image-1B-Base',
+        'identifier': ['1b', '2b'],
+        'target_subfolder': 'Qwen3-VL-2B-Instruct',
+    },
+    'Qwen3-VL 4B Conditional': {
+        'cls': transformers.Qwen3VLForConditionalGeneration,
+        'target_repo': 'SeFi-Image/SeFi-Image-5B-Base',
+        'identifier': ['5b'],
+        'target_subfolder': 'Qwen3-VL-4B-Instruct',
+    },
     'Qwen3-VL 8B Conditional': {
         'cls': transformers.Qwen3VLForConditionalGeneration,
         'target_repo': 'Boogu/Boogu-Image-0.1-Base',

pipelines/generic_text_encoder.py

@@ -16,7 +16,13 @@ def get_shared(cls, repo_id, subfolder=None, variant=None):
     if variant is not None:
         args['variant'] = variant
     for name, item in shared_te_map.items():
-        if item['cls'] == cls and (item.get('identifier', None) is None or item.get('identifier', None).lower() in repo_id.lower()):
+        identifiers = item.get('identifier', [])
+        if identifiers is None:
+            identifiers = []
+        if isinstance(identifiers, str):
+            identifiers = [identifiers]
+        identifiers = [identifier.lower() for identifier in identifiers if identifier is not None]
+        if item['cls'] == cls and (not identifiers or any(identifier in repo_id.lower() for identifier in identifiers)):
             if item.get('config_class', None) is not None and item.get('config_path', None) is not None:
                 with open(item['config_path'], encoding='utf8') as f:
                     args['config'] = item['config_class'](**json.load(f))

pipelines/krea2/krea2_lora.py

@@ -4,10 +4,15 @@ Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
 (``lora_force_diffusers`` off and ``krea2`` in ``allow_native``).
 
 The transformer module tree mirrors the checkpoint (``blocks.N.attn.{wq,wk,wv,wo,gate}``,
-``blocks.N.mlp.{gate,up,down}``, ``txtfusion.*``, ``first``, ``last`` ...), so dotted keys
-bind verbatim with no name rewrite and no fused-QKV split. Kohya flat-underscore keys are
-reconstructed back to dotted paths, protecting the two compound module names
-(``layerwise_blocks``, ``refiner_blocks``).
+``blocks.N.mlp.{gate,up,down}``, ``txtfusion.*``, ``first``, ``last`` ...), so
+checkpoint-style dotted keys bind verbatim with no fused-QKV split. Upstream diffusers
+(``Krea2Transformer2DModel``) names the same modules differently
+(``transformer_blocks.N.attn.to_q``, ``text_fusion.*``, ``img_in``, ``txt_in``,
+``time_embed``, ``time_mod_proj``, ``final_layer``); LoRAs trained against it (e.g. the
+official ``krea/Krea-2-LoRA-*`` releases) are rewritten to checkpoint names via the
+``DIFFUSERS_*_MAP`` tables. Kohya flat-underscore keys are reconstructed back to dotted
+checkpoint paths, protecting the two compound module names (``layerwise_blocks``,
+``refiner_blocks``).
 """
 
 from modules.lora import native_adapter
@@ -15,24 +20,122 @@ from modules.lora import native_adapter
 
 KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
 
-# Top-level module names that a bare diffusers-format LoRA key can start with.
-BARE_DIFFUSERS_PREFIXES = ("blocks.", "txtfusion.", "first.", "last.", "tmlp.", "tproj.", "txtmlp.")
+# Top-level module names that a bare LoRA key can start with: the transformer's own
+# checkpoint-style names plus the upstream-diffusers names (as saved by
+# ``Krea2Transformer2DModel.save_lora_adapter()``).
+BARE_DIFFUSERS_PREFIXES = (
+    "blocks.", "txtfusion.", "first.", "last.", "tmlp.", "tproj.", "txtmlp.",
+    "transformer_blocks.", "text_fusion.", "img_in.", "txt_in.", "time_embed.", "time_mod_proj.", "final_layer.",
+)
+
+# Upstream-diffusers attention/ff leaves -> checkpoint leaves (block-level modules).
+DIFFUSERS_LEAF_MAP = {
+    ".attn.to_q": ".attn.wq",
+    ".attn.to_k": ".attn.wk",
+    ".attn.to_v": ".attn.wv",
+    ".attn.to_gate": ".attn.gate",
+    ".attn.to_out.0": ".attn.wo",
+}
+
+# Upstream-diffusers non-block modules -> checkpoint paths (exact match on the full base).
+# Sequential containers on the checkpoint side are addressed by index
+# (``tmlp``/``txtmlp``/``tproj`` interleave activations and norms with the Linears).
+DIFFUSERS_EXTRA_MAP = {
+    "img_in": "first",
+    "txt_in.linear_1": "txtmlp.1",
+    "txt_in.linear_2": "txtmlp.3",
+    "time_embed.linear_1": "tmlp.0",
+    "time_embed.linear_2": "tmlp.2",
+    "time_mod_proj": "tproj.1",
+    "final_layer.linear": "last.linear",
+}
+
+
+# === Re-exports for test/back-compat ===
+# The offline test suite addresses the family suffix/marker tables and the
+# parse helpers through this module surface rather than importing native_adapter.
+
+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
+has_marker = native_adapter.has_marker
+
+
+def parse_key(key, suffixes):
+    """Krea2-bound :func:`native_adapter.parse_key`."""
+    return native_adapter.parse_key(
+        key, suffixes,
+        prefixes=KNOWN_PREFIXES,
+        bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES,
+    )
+
+
+def group_by_suffixes(state_dict, suffixes):
+    """Krea2-bound :func:`native_adapter.group_by_suffixes`."""
+    return native_adapter.group_by_suffixes(
+        state_dict, suffixes,
+        prefixes=KNOWN_PREFIXES,
+        bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES,
+    )
 
 
 def resolve_targets(prefix_used, base):
-    """Return ``[(diffusers_path, None), ...]`` for a parsed group key.
+    """Return ``[(checkpoint_path, None), ...]`` for a parsed group key.
 
-    K2's diffusers module names equal the checkpoint names, so dotted keys map verbatim.
-    Universal passthrough prefixes are handled upstream by
+    Upstream-diffusers paths are rewritten to checkpoint names; checkpoint-style
+    paths bind verbatim. For the passthrough prefixes (``transformer.`` and the
+    bare-diffusers sentinel) returning ``[]`` defers to the verbatim fallback in
     :func:`native_adapter.resolve_group_targets`.
     """
-    if prefix_used in (None, "diffusion_model.", "transformer."):
-        return [(base, None)]
-    if prefix_used in ("lora_unet_", "lora_transformer_"):
+    if prefix_used in (None, "diffusion_model."):
+        return _diffusers_to_checkpoint(base) or [(base, None)]
+    if prefix_used in ("transformer.", native_adapter.BARE_DIFFUSERS_PREFIX_USED):
+        return _diffusers_to_checkpoint(base)
+    if prefix_used == "lora_unet_":
         return _underscore_to_dotted(base)
     return []
 
 
+def _diffusers_to_checkpoint(base):
+    """Rewrite an upstream-diffusers module path to the checkpoint-style path.
+
+    Returns ``[]`` for paths not in upstream-diffusers form so callers (and the
+    shared verbatim fallback) handle checkpoint-style keys.
+    """
+    extra = DIFFUSERS_EXTRA_MAP.get(base)
+    if extra is not None:
+        return [(extra, None)]
+    if base.startswith("transformer_blocks."):
+        path = "blocks." + base[len("transformer_blocks."):]
+    elif base.startswith("text_fusion."):
+        path = "txtfusion." + base[len("text_fusion."):]
+    else:
+        return []
+    for leaf, renamed in DIFFUSERS_LEAF_MAP.items():
+        if path.endswith(leaf):
+            path = path[:-len(leaf)] + renamed
+            break
+    path = path.replace(".ff.", ".mlp.")
+    return [(path, None)]
+
+
 def _underscore_to_dotted(base):
     """Rebuild a dotted path from a kohya flat-underscore base, keeping compound names intact."""
     protected = base.replace("layerwise_blocks", "layerwise@blocks").replace("refiner_blocks", "refiner@blocks")

pipelines/model_sefi.py

@@ -0,0 +1,42 @@
+import transformers
+import diffusers
+from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_sefi(checkpoint_info, diffusers_load_config=None):
+    if diffusers_load_config is None:
+        diffusers_load_config = {}
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+
+    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config)
+    log.debug(f'Load model: type=SeFi repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+    from pipelines.sefi import SeFiTransformer2DModel, SeFiPipeline
+
+    transformer = generic.load_transformer(repo_id, cls_name=SeFiTransformer2DModel, load_config=diffusers_load_config)
+    text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLForConditionalGeneration, load_config=diffusers_load_config)
+    if repo_id is None or repo_id.lower() == 'none':
+        return None
+
+    pipe = SeFiPipeline.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_args,
+    )
+
+    diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sefi"] = SeFiPipeline
+
+    generic.load_vae_override(pipe, diffusers_load_config)
+
+    del text_encoder
+    del transformer
+    sd_hijack_te.init_hijack(pipe)
+    sd_hijack_vae.init_hijack(pipe)
+    devices.torch_gc(force=True, reason='load')
+
+    return pipe

pipelines/native_transformer.py

@@ -867,7 +867,7 @@ def build_component_prequantized(
         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,
+        use_quantized_matmul=(shared.opts.sdnq_quantize_matmul_mode != "disabled"),
         dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
         add_skip_keys=False,
         modules_to_not_convert=[],
@@ -976,7 +976,7 @@ def build_component_prequantized(
         component,
         dtype=target_dtype,
         dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
-        use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+        use_quantized_matmul=(shared.opts.sdnq_quantize_matmul_mode != "disabled"),
     )
     return component
 

pipelines/sefi/__init__.py

@@ -0,0 +1,3 @@
+from .transformer_sefi import SeFiTransformer2DModel
+from .pipeline_sefi import SeFiPipeline
+from .pipeline_output import SeFiPipelineOutput

pipelines/sefi/pipeline_output.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+import numpy as np
+import PIL.Image
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SeFiPipelineOutput(BaseOutput):
+    """
+    Output class for SeFi-Image pipelines.
+
+    Args:
+        images (`list[PIL.Image.Image]` or `np.ndarray`)
+            Generated images.
+    """
+
+    images: list[PIL.Image.Image] | np.ndarray

pipelines/sefi/pipeline_sefi.py

@@ -0,0 +1,700 @@
+# Copyright 2026 SeFi-Image Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable
+
+import torch
+from transformers import Qwen2Tokenizer, Qwen3VLForConditionalGeneration
+
+from diffusers.models import AutoencoderKL, AutoencoderKLFlux2
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from .pipeline_output import SeFiPipelineOutput
+from .transformer_sefi import SeFiTransformer2DModel
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import SeFiPipeline
+
+        >>> pipe = SeFiPipeline.from_pretrained("./sefi-1b-base-diffusers", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> image = pipe("A red apple on a wooden table.").images[0]
+        >>> image.save("sefi.png")
+        ```
+"""
+
+
+SUPPORTED_TURBO_STEPS = {4, 8, 10}
+
+
+def _apply_timestep_shift_unit_interval(u_unit: torch.Tensor, alpha: float) -> torch.Tensor:
+    alpha = float(alpha)
+    if alpha <= 0:
+        raise ValueError(f"`timestep_shift_alpha` must be > 0, got {alpha}.")
+    if alpha == 1.0:
+        return u_unit
+    denominator = 1.0 + (alpha - 1.0) * u_unit
+    return (alpha * u_unit) / denominator
+
+
+def _combine_guided_velocity(base_pred: torch.Tensor, cond_pred: torch.Tensor, guidance_scale: float) -> torch.Tensor:
+    return base_pred + float(guidance_scale) * (cond_pred - base_pred)
+
+
+class SeFiPipeline(DiffusionPipeline):
+    r"""
+    SeFi-Image text-to-image generation pipeline.
+
+    Args:
+        transformer ([`SeFiTransformer2DModel`]):
+            Transformer that predicts semantic and texture latent velocities.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Flow-matching scheduler whose training timesteps and sigmas are used for SeFi's dual-time update.
+        vae ([`AutoencoderKL`] or [`AutoencoderKLFlux2`]):
+            Texture VAE used to decode the final texture latent stream.
+        text_encoder ([`~transformers.Qwen3VLForConditionalGeneration`]):
+            Qwen3-VL text encoder. SeFi uses concatenated hidden states from selected text layers.
+        tokenizer ([`~transformers.Qwen2Tokenizer`]):
+            Tokenizer paired with the Qwen3-VL text encoder.
+        semantic_channels (`int`, defaults to `16`):
+            Number of semantic latent channels.
+        texture_vae_name (`str`, defaults to `"flux2"`):
+            Texture VAE normalization type. Supported values are `"sd1.5"`, `"flux1"`, and `"flux2"`.
+        is_turbo (`bool`, defaults to `False`):
+            Whether the checkpoint is a distilled Turbo model.
+        default_guidance_scale (`float`, defaults to `4.0`):
+            Default guidance scale used when `guidance_scale` is not provided.
+        default_num_inference_steps (`int`, defaults to `50`):
+            Default number of inference steps used when `num_inference_steps` is not provided.
+        delta_t (`float`, defaults to `0.1`):
+            Semantic stream lead over the texture stream.
+        timestep_shift_alpha (`float`, defaults to `0.3`):
+            Unit-interval timestep shift applied before the SeFi dual-time schedule.
+        text_encoder_hidden_layers (`tuple[int, ...]`, defaults to `(9, 18, 27)`):
+            Text encoder hidden-state indices concatenated as prompt embeddings.
+        max_sequence_length (`int`, defaults to `1024`):
+            Maximum prompt token length.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        transformer: SeFiTransformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL | AutoencoderKLFlux2,
+        text_encoder: Qwen3VLForConditionalGeneration,
+        tokenizer: Qwen2Tokenizer,
+        semantic_channels: int = 16,
+        texture_vae_name: str = "flux2",
+        is_turbo: bool = False,
+        default_guidance_scale: float = 4.0,
+        default_num_inference_steps: int = 50,
+        delta_t: float = 0.1,
+        timestep_shift_alpha: float = 0.3,
+        text_encoder_hidden_layers: list[int] | tuple[int, ...] = (9, 18, 27),
+        max_sequence_length: int = 1024,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            transformer=transformer,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )
+        if isinstance(text_encoder_hidden_layers, str):
+            text_encoder_hidden_layers = tuple(int(layer) for layer in text_encoder_hidden_layers.split(","))
+        semantic_channels = 16 if semantic_channels is None else semantic_channels
+        if texture_vae_name is None:
+            texture_vae_name = "flux2" if vae is not None and hasattr(vae, "bn") else "sd1.5"
+        default_guidance_scale = 4.0 if default_guidance_scale is None else default_guidance_scale
+        default_num_inference_steps = 50 if default_num_inference_steps is None else default_num_inference_steps
+        text_encoder_hidden_layers = (9, 18, 27) if text_encoder_hidden_layers is None else text_encoder_hidden_layers
+        max_sequence_length = 1024 if max_sequence_length is None else max_sequence_length
+        self.register_to_config(
+            semantic_channels=semantic_channels,
+            texture_vae_name=texture_vae_name,
+            is_turbo=is_turbo,
+            default_guidance_scale=default_guidance_scale,
+            default_num_inference_steps=default_num_inference_steps,
+            delta_t=delta_t,
+            timestep_shift_alpha=timestep_shift_alpha,
+            text_encoder_hidden_layers=tuple(text_encoder_hidden_layers),
+            max_sequence_length=max_sequence_length,
+        )
+
+        self.semantic_channels = int(semantic_channels)
+        self.texture_vae_name = str(texture_vae_name).lower()
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.default_sample_size = 128
+        self._guidance_scale = None
+        self._attention_kwargs = None
+        self._current_timestep = None
+        self._interrupt = False
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self.guidance_scale is not None and self.guidance_scale > 1.0
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @staticmethod
+    def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
+        B, L, _ = x.shape
+        out_ids = []
+
+        for i in range(B):
+            t = torch.arange(1) if t_coord is None else t_coord[i]
+            h = torch.arange(1)
+            w = torch.arange(1)
+            l = torch.arange(L)
+
+            coords = torch.cartesian_prod(t, h, w, l)
+            out_ids.append(coords)
+
+        return torch.stack(out_ids)
+
+    @staticmethod
+    def _prepare_latent_ids(latents: torch.Tensor):
+        r"""
+        Generates 4D position coordinates (T, H, W, L) for latent tensors.
+
+        Args:
+            latents (torch.Tensor):
+                Latent tensor of shape (B, C, H, W)
+
+        Returns:
+            torch.Tensor:
+                Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0,
+                H=[0..H-1], W=[0..W-1], L=0
+        """
+
+        batch_size, _, height, width = latents.shape
+
+        t = torch.arange(1)  # [0] - time dimension
+        h = torch.arange(height)
+        w = torch.arange(width)
+        l = torch.arange(1)  # [0] - layer dimension
+
+        # Create position IDs: (H*W, 4)
+        latent_ids = torch.cartesian_prod(t, h, w, l)
+
+        # Expand to batch: (B, H*W, 4)
+        latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+        return latent_ids
+
+    @staticmethod
+    def _unpatchify_latents(latents):
+        batch_size, num_channels_latents, height, width = latents.shape
+        latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
+        return latents
+
+    @staticmethod
+    def _pack_latents(latents):
+        """
+        pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)
+        """
+
+        batch_size, num_channels, height, width = latents.shape
+        latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents_with_ids(
+        x: torch.Tensor, x_ids: torch.Tensor, height: int | None = None, width: int | None = None
+    ):
+        """
+        using position ids to scatter tokens into place
+        """
+        x_list = []
+        for data, pos in zip(x, x_ids):
+            _, ch = data.shape
+            h_ids = pos[:, 1].to(torch.int64)
+            w_ids = pos[:, 2].to(torch.int64)
+
+            h = torch.max(h_ids) + 1
+            w = torch.max(w_ids) + 1
+
+            flat_ids = h_ids * w + w_ids
+
+            out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype)
+            out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data)
+
+            # reshape from (H * W, C) to (H, W, C) and permute to (C, H, W)
+
+            out = out.view(h, w, ch).permute(2, 0, 1)
+            x_list.append(out)
+
+        return torch.stack(x_list, dim=0)
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        if height is not None and height <= 0:
+            raise ValueError(f"`height` must be > 0, got {height}.")
+        if width is not None and width <= 0:
+            raise ValueError(f"`width` must be > 0, got {width}.")
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`, not both.")
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
+        if negative_prompt_embeds is not None and prompt_embeds is None:
+            raise ValueError("`negative_prompt_embeds` requires `prompt_embeds`.")
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found "
+                f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+    def _build_chat_text(self, prompt: str) -> str:
+        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+        try:
+            return self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+        except TypeError:
+            return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    def _align_text_encoder_rotary_dtype(self, device: torch.device):
+        text_encoder = self.text_encoder
+        if text_encoder is None:
+            return
+
+        try:
+            text_encoder_dtype = next(text_encoder.parameters()).dtype
+        except StopIteration:
+            return
+
+        text_model = text_encoder.model if hasattr(text_encoder, "model") else text_encoder
+        language_model = getattr(text_model, "language_model", None)
+        rotary_emb = getattr(language_model, "rotary_emb", None)
+        if rotary_emb is not None:
+            # Qwen3-VL stores RoPE inverse frequencies as non-persistent buffers. `from_pretrained(torch_dtype=...)`
+            # can leave them in fp32 even when text weights are bf16, while the reference SeFi wrapper casts the whole
+            # text encoder module. Keep these buffers aligned before text encoding.
+            rotary_emb.to(device=device, dtype=text_encoder_dtype)
+
+    def _get_qwen3vl_prompt_embeds(
+        self,
+        prompt: str | list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        max_sequence_length: int,
+        hidden_layers: tuple[int, ...],
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        chat_texts = [self._build_chat_text(single_prompt) for single_prompt in prompt]
+        tokenized = self.tokenizer(
+            chat_texts,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=max_sequence_length,
+        )
+        input_ids = tokenized["input_ids"].to(device)
+        attention_mask = tokenized["attention_mask"].to(device)
+
+        self._align_text_encoder_rotary_dtype(device)
+        outputs = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+            logits_to_keep=1,
+            return_dict=True,
+        )
+        hidden_states = outputs.hidden_states
+        max_idx = len(hidden_states) - 1
+        for layer_idx in hidden_layers:
+            if layer_idx > max_idx:
+                raise ValueError(
+                    f"Requested hidden layer {layer_idx}, but text encoder only provides up to {max_idx}."
+                )
+
+        stacked = torch.stack([hidden_states[idx] for idx in hidden_layers], dim=1)
+        stacked = stacked.to(dtype=dtype, device=device)
+        batch_size, num_layers, seq_len, hidden_dim = stacked.shape
+        prompt_embeds = stacked.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_layers * hidden_dim)
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        max_sequence_length: int | None = None,
+        text_encoder_hidden_layers: tuple[int, ...] | None = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or (self.transformer.dtype if self.transformer is not None else self.text_encoder.dtype)
+        max_sequence_length = max_sequence_length or self.config.max_sequence_length
+        text_encoder_hidden_layers = text_encoder_hidden_layers or tuple(self.config.text_encoder_hidden_layers)
+
+        if prompt_embeds is None:
+            prompt_embeds = self._get_qwen3vl_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                dtype=dtype,
+                max_sequence_length=max_sequence_length,
+                hidden_layers=tuple(text_encoder_hidden_layers),
+            )
+        else:
+            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+
+        batch_size, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        text_ids = self._prepare_text_ids(prompt_embeds).to(device)
+        return prompt_embeds, text_ids
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ):
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+        shape = (batch_size, self.transformer.config.in_channels, height // 2, width // 2)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+            if tuple(latents.shape) != tuple(shape):
+                raise ValueError(f"Unexpected `latents` shape {tuple(latents.shape)}, expected {tuple(shape)}.")
+
+        latent_ids = self._prepare_latent_ids(latents).to(device)
+        return latents, latent_ids
+
+    def _timesteps_and_sigmas(self, u_continuous: torch.Tensor, n_dim: int, dtype: torch.dtype):
+        num_steps = int(self.scheduler.config.num_train_timesteps)
+        indices = (u_continuous * (num_steps - 1)).long().clamp(0, num_steps - 1)
+        timesteps = self.scheduler.timesteps[indices.cpu()].to(self._execution_device)
+        sigmas = self.scheduler.sigmas[indices.cpu()].to(device=self._execution_device, dtype=dtype)
+        while sigmas.ndim < n_dim:
+            sigmas = sigmas.unsqueeze(-1)
+        return timesteps, sigmas
+
+    def decode_texture_latents(self, texture_latents: torch.Tensor, output_type: str = "pil"):
+        if self.texture_vae_name == "flux2":
+            if not hasattr(self.vae, "bn"):
+                raise ValueError("`texture_vae_name='flux2'` requires a VAE with batch-norm statistics.")
+            eps = float(getattr(self.vae.config, "batch_norm_eps", 1e-6))
+            bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(texture_latents.device, texture_latents.dtype)
+            bn_std = torch.sqrt(
+                self.vae.bn.running_var.view(1, -1, 1, 1).to(texture_latents.device, texture_latents.dtype) + eps
+            )
+            texture_latents = texture_latents * bn_std + bn_mean
+            raw_latents = self._unpatchify_latents(texture_latents)
+        else:
+            scaling_factor = float(getattr(self.vae.config, "scaling_factor", 1.0))
+            shift_factor = float(getattr(self.vae.config, "shift_factor", 0.0) or 0.0)
+            raw_latents = self._unpatchify_latents(texture_latents)
+            raw_latents = raw_latents / scaling_factor + shift_factor
+
+        image = self.vae.decode(raw_latents.to(dtype=self.vae.dtype), return_dict=False)[0]
+        return self.image_processor.postprocess(image, output_type=output_type)
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict | None = None,
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int | None = None,
+        text_encoder_hidden_layers: tuple[int, ...] | None = None,
+    ) -> SeFiPipelineOutput | tuple:
+        r"""
+        Generates images from text prompts with SeFi-Image.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to guide image generation.
+            height (`int`, *optional*):
+                Height in pixels of the generated image.
+            width (`int`, *optional*):
+                Width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*):
+                Number of denoising steps. Base/RL checkpoints default to 50 and Turbo checkpoints default to 4.
+            guidance_scale (`float`, *optional*):
+                Classifier-free guidance scale. Turbo checkpoints require `guidance_scale=1.0`.
+            num_images_per_prompt (`int`, defaults to `1`):
+                Number of images to generate per prompt.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Random generator for deterministic generation.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated semantic and texture latents.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated prompt embeddings.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative prompt embeddings.
+            output_type (`str`, defaults to `"pil"`):
+                Output type of the generated image. Choose between `"pil"`, `"np"`, and `"latent"`.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a [`SeFiPipelineOutput`] instead of a tuple.
+            attention_kwargs (`dict`, *optional*):
+                Keyword arguments passed to attention processors.
+            callback_on_step_end (`Callable`, *optional*):
+                Function called at the end of each denoising step.
+            callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`):
+                Tensor inputs passed to `callback_on_step_end`.
+            max_sequence_length (`int`, *optional*):
+                Maximum prompt sequence length.
+            text_encoder_hidden_layers (`tuple[int, ...]`, *optional*):
+                Text encoder hidden-state layers to concatenate.
+
+        Examples:
+
+        Returns:
+            [`SeFiPipelineOutput`] or `tuple`: Generated images.
+        """
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = int(num_inference_steps or self.config.default_num_inference_steps)
+        guidance_scale = float(guidance_scale if guidance_scale is not None else self.config.default_guidance_scale)
+
+        self.check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        if self.config.is_turbo:
+            if num_inference_steps not in SUPPORTED_TURBO_STEPS:
+                raise ValueError(f"SeFi Turbo models support {sorted(SUPPORTED_TURBO_STEPS)} steps.")
+            if guidance_scale != 1.0:
+                raise ValueError("SeFi Turbo models should run with `guidance_scale=1.0`.")
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        prompt_embeds, text_ids = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            device=device,
+            dtype=dtype,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            text_encoder_hidden_layers=text_encoder_hidden_layers,
+        )
+
+        negative_text_ids = None
+        if self.do_classifier_free_guidance:
+            if negative_prompt_embeds is None:
+                negative_prompt = "" if batch_size == 1 else [""] * batch_size
+            else:
+                negative_prompt = None
+            negative_prompt_embeds, negative_text_ids = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                device=device,
+                dtype=dtype,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                text_encoder_hidden_layers=text_encoder_hidden_layers,
+            )
+
+        latents, latent_ids = self.prepare_latents(
+            batch_size=batch_size * num_images_per_prompt,
+            height=height,
+            width=width,
+            dtype=dtype,
+            device=device,
+            generator=generator,
+            latents=latents,
+        )
+
+        u_base_unit = torch.linspace(
+            0.0,
+            1.0,
+            steps=num_inference_steps + 1,
+            device=device,
+            dtype=torch.float32,
+        )
+        u_shifted_unit = _apply_timestep_shift_unit_interval(u_base_unit, self.config.timestep_shift_alpha)
+        _, base_sigmas_schedule = self._timesteps_and_sigmas(u_shifted_unit, n_dim=1, dtype=torch.float32)
+        u_sem_raw_schedule = u_shifted_unit * (1.0 + float(self.config.delta_t))
+
+        self._num_timesteps = num_inference_steps
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i in range(num_inference_steps):
+                if self.interrupt:
+                    continue
+
+                u_sem_raw_cur = torch.full((latents.shape[0],), float(u_sem_raw_schedule[i].item()), device=device)
+                u_sem_raw_next = torch.full(
+                    (latents.shape[0],), float(u_sem_raw_schedule[i + 1].item()), device=device
+                )
+                u_tex_cur = torch.clamp(u_sem_raw_cur - float(self.config.delta_t), min=0.0, max=1.0)
+                u_sem_cur = torch.clamp(u_sem_raw_cur, max=1.0)
+                u_tex_next = torch.clamp(u_sem_raw_next - float(self.config.delta_t), min=0.0, max=1.0)
+                u_sem_next = torch.clamp(u_sem_raw_next, max=1.0)
+
+                timesteps_sem_cur, sigmas_sem_cur = self._timesteps_and_sigmas(u_sem_cur, latents.ndim, latents.dtype)
+                timesteps_tex_cur, sigmas_tex_cur = self._timesteps_and_sigmas(u_tex_cur, latents.ndim, latents.dtype)
+                _, sigmas_sem_next = self._timesteps_and_sigmas(u_sem_next, latents.ndim, latents.dtype)
+                _, sigmas_tex_next = self._timesteps_and_sigmas(u_tex_next, latents.ndim, latents.dtype)
+
+                self._current_timestep = base_sigmas_schedule[i]
+                packed_latents = self._pack_latents(latents)
+                pred_cond = self.transformer(
+                    hidden_states=packed_latents,
+                    timestep_sem=timesteps_sem_cur / 1000,
+                    timestep_tex=timesteps_tex_cur / 1000,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_ids,
+                    joint_attention_kwargs=self.attention_kwargs,
+                    return_dict=False,
+                )[0]
+                pred_cond = pred_cond[:, : packed_latents.size(1)]
+                pred_cond = self._unpack_latents_with_ids(pred_cond, latent_ids)
+
+                if self.do_classifier_free_guidance:
+                    pred_uncond = self.transformer(
+                        hidden_states=packed_latents,
+                        timestep_sem=timesteps_sem_cur / 1000,
+                        timestep_tex=timesteps_tex_cur / 1000,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=negative_text_ids,
+                        img_ids=latent_ids,
+                        joint_attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    pred_uncond = pred_uncond[:, : packed_latents.size(1)]
+                    pred_uncond = self._unpack_latents_with_ids(pred_uncond, latent_ids)
+                    velocity = _combine_guided_velocity(pred_uncond, pred_cond, guidance_scale)
+                else:
+                    velocity = pred_cond
+
+                vel_sem = velocity[:, : self.semantic_channels]
+                vel_tex = velocity[:, self.semantic_channels :]
+                lat_sem = latents[:, : self.semantic_channels]
+                lat_tex = latents[:, self.semantic_channels :]
+
+                lat_sem = lat_sem + (sigmas_sem_next - sigmas_sem_cur) * vel_sem
+                lat_tex = lat_tex + (sigmas_tex_next - sigmas_tex_cur) * vel_tex
+                latents = torch.cat([lat_sem, lat_tex], dim=1)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, self._current_timestep, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+                progress_bar.update()
+
+        if output_type == "latent":
+            image = latents
+        else:
+            texture_latents = latents[:, self.semantic_channels :]
+            image = self.decode_texture_latents(texture_latents, output_type=output_type)
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return SeFiPipelineOutput(images=image)

pipelines/sefi/transformer_sefi.py

@@ -0,0 +1,282 @@
+# Copyright 2026 SeFi-Image Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, apply_lora_scale
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+
+
+@dataclass
+class SeFiTransformer2DModelOutput(BaseOutput):
+    """
+    Output of [`SeFiTransformer2DModel`].
+
+    Args:
+        sample (`torch.Tensor` of shape `(batch_size, image_sequence_length, out_channels)`):
+            Predicted velocity for packed semantic and texture latents.
+    """
+
+    sample: torch.Tensor
+
+
+class SeFiDualTimestepEmbeddings(nn.Module):
+    """Dual semantic/texture timestep embedding used by SeFi-Image."""
+
+    def __init__(self, in_channels: int, embedding_dim: int, bias: bool = False):
+        super().__init__()
+        if embedding_dim % 2 != 0:
+            raise ValueError(f"`embedding_dim` must be even for dual timestep embeddings, got {embedding_dim}.")
+
+        half_dim = embedding_dim // 2
+        self.time_proj = Timesteps(
+            num_channels=int(in_channels),
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.semantic_embedder = TimestepEmbedding(
+            in_channels=int(in_channels),
+            time_embed_dim=half_dim,
+            sample_proj_bias=bias,
+        )
+        self.texture_embedder = TimestepEmbedding(
+            in_channels=int(in_channels),
+            time_embed_dim=half_dim,
+            sample_proj_bias=bias,
+        )
+
+    def forward(self, timestep_sem: torch.Tensor, timestep_tex: torch.Tensor) -> torch.Tensor:
+        sem_proj = self.time_proj(timestep_sem)
+        tex_proj = self.time_proj(timestep_tex)
+        sem_emb = self.semantic_embedder(sem_proj.to(timestep_sem.dtype))
+        tex_emb = self.texture_embedder(tex_proj.to(timestep_tex.dtype))
+        return torch.cat([sem_emb, tex_emb], dim=-1)
+
+
+class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
+    """
+    SeFi-Image transformer with explicit semantic and texture timestep conditioning.
+
+    SeFi-Image reuses a Flux2-style MMDiT backbone, but replaces the single timestep embedding with a dual embedding:
+    one timestep for the semantic latent stream and one timestep for the texture latent stream.
+
+    Args:
+        patch_size (`int`, defaults to `1`):
+            Patch size of the Flux2 backbone.
+        in_channels (`int`, defaults to `128`):
+            Number of packed latent channels. This is `semantic_channels + texture_channels`.
+        out_channels (`int`, *optional*):
+            Number of output packed latent channels. Defaults to `in_channels`.
+        num_layers (`int`, defaults to `4`):
+            Number of double-stream transformer layers.
+        num_single_layers (`int`, defaults to `12`):
+            Number of single-stream transformer layers.
+        attention_head_dim (`int`, defaults to `128`):
+            Dimension per attention head.
+        num_attention_heads (`int`, defaults to `16`):
+            Number of attention heads.
+        joint_attention_dim (`int`, defaults to `6144`):
+            Dimension of the concatenated Qwen3-VL hidden states.
+        timestep_guidance_channels (`int`, defaults to `256`):
+            Number of channels for sinusoidal timestep projection.
+        mlp_ratio (`float`, defaults to `3.0`):
+            MLP expansion ratio in transformer blocks.
+        axes_dims_rope (`tuple[int, ...]`, defaults to `(32, 32, 32, 32)`):
+            RoPE dimensions for Flux2 positional embeddings.
+        rope_theta (`int`, defaults to `2000`):
+            RoPE theta.
+        eps (`float`, defaults to `1e-6`):
+            Normalization epsilon.
+        text_input_dim (`int`, *optional*):
+            Expected text embedding dimension. Defaults to `joint_attention_dim`.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 128,
+        out_channels: int | None = None,
+        num_layers: int = 4,
+        num_single_layers: int = 12,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 16,
+        joint_attention_dim: int = 6144,
+        timestep_guidance_channels: int = 256,
+        mlp_ratio: float = 3.0,
+        axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
+        rope_theta: int = 2000,
+        eps: float = 1e-6,
+        text_input_dim: int | None = None,
+    ):
+        super().__init__()
+
+        text_input_dim = joint_attention_dim if text_input_dim is None else text_input_dim
+        if int(text_input_dim) != int(joint_attention_dim):
+            raise ValueError(
+                f"`text_input_dim` must match `joint_attention_dim`, got {text_input_dim} and {joint_attention_dim}."
+            )
+
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.backbone = Flux2Transformer2DModel(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            num_single_layers=num_single_layers,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            joint_attention_dim=joint_attention_dim,
+            timestep_guidance_channels=timestep_guidance_channels,
+            mlp_ratio=mlp_ratio,
+            axes_dims_rope=axes_dims_rope,
+            rope_theta=rope_theta,
+            eps=eps,
+            guidance_embeds=False,
+        )
+        # The reference SeFi transformer deletes Flux2's timestep/guidance embedder and stores only the dual embedder.
+        self.backbone.time_guidance_embed = nn.Identity()
+        self.dual_time_embed = SeFiDualTimestepEmbeddings(
+            in_channels=timestep_guidance_channels,
+            embedding_dim=self.inner_dim,
+            bias=False,
+        )
+        self.gradient_checkpointing = False
+
+    @apply_lora_scale("joint_attention_kwargs")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep_sem: torch.Tensor,
+        timestep_tex: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        txt_ids: torch.Tensor,
+        img_ids: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | SeFiTransformer2DModelOutput:
+        """
+        The [`SeFiTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Packed semantic and texture latents of shape `(batch_size, image_sequence_length, in_channels)`.
+            timestep_sem (`torch.Tensor`):
+                Semantic stream timesteps, normalized to the Diffusers convention where `1.0` corresponds to `1000`.
+            timestep_tex (`torch.Tensor`):
+                Texture stream timesteps, normalized to the Diffusers convention where `1.0` corresponds to `1000`.
+            encoder_hidden_states (`torch.Tensor`):
+                Text conditioning embeddings.
+            txt_ids (`torch.Tensor`):
+                Text token position ids.
+            img_ids (`torch.Tensor`):
+                Image token position ids.
+            joint_attention_kwargs (`dict`, *optional*):
+                Keyword arguments forwarded to attention processors.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return [`SeFiTransformer2DModelOutput`] or a tuple.
+
+        Returns:
+            [`SeFiTransformer2DModelOutput`] or `tuple`:
+                Predicted semantic and texture latent velocities.
+        """
+
+        num_txt_tokens = encoder_hidden_states.shape[1]
+
+        timestep_sem = timestep_sem.to(hidden_states.dtype) * 1000
+        timestep_tex = timestep_tex.to(hidden_states.dtype) * 1000
+        temb = self.dual_time_embed(timestep_sem, timestep_tex)
+
+        double_stream_mod_img = self.backbone.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.backbone.double_stream_modulation_txt(temb)
+        single_stream_mod = self.backbone.single_stream_modulation(temb)
+
+        hidden_states = self.backbone.x_embedder(hidden_states)
+        encoder_hidden_states = self.backbone.context_embedder(encoder_hidden_states)
+
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        image_rotary_emb = self.backbone.pos_embed(img_ids)
+        text_rotary_emb = self.backbone.pos_embed(txt_ids)
+        concat_rotary_emb = (
+            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+        )
+
+        for block in self.backbone.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    double_stream_mod_img,
+                    double_stream_mod_txt,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb_mod_img=double_stream_mod_img,
+                    temb_mod_txt=double_stream_mod_txt,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for block in self.backbone.single_transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    None,
+                    single_stream_mod,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=None,
+                    temb_mod=single_stream_mod,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.backbone.norm_out(hidden_states, temb)
+        output = self.backbone.proj_out(hidden_states)
+
+        if not return_dict:
+            return (output,)
+
+        return SeFiTransformer2DModelOutput(sample=output)

pyproject.toml

@@ -130,33 +130,34 @@ main.ignore-paths=[
     "modules/teacache",
     "modules/todo",
     "modules/res4lyf",
+    "pipelines/boogu",
     "pipelines/bria",
-    "pipelines/flex2",
+    "pipelines/chrono",
     "pipelines/f_lite",
-    "pipelines/hidream",
+    "pipelines/flex2",
     "pipelines/hdm",
+    "pipelines/hidream",
+    "pipelines/lumina_dimmo",
     "pipelines/meissonic",
     "pipelines/omnigen2",
     "pipelines/segmoe",
-    "pipelines/xomni",
-    "pipelines/chrono",
+    "pipelines/sefi",
     "pipelines/step1x",
-    "pipelines/vibe",
     "pipelines/ultraflux",
-    "pipelines/lumina_dimmo",
-    "pipelines/boogu",
+    "pipelines/vibe",
+    "pipelines/xomni",
     "scripts/consistory",
     "scripts/ctrlx",
     "scripts/daam",
     "scripts/demofusion",
+    "scripts/differential_diffusion.py",
     "scripts/freescale",
     "scripts/infiniteyou",
     "scripts/instantir",
-    "scripts/lbm",
     "scripts/layerdiffuse",
+    "scripts/lbm",
     "scripts/mod",
     "scripts/pixelsmith",
-    "scripts/differential_diffusion.py",
     "scripts/pulid",
     "scripts/xadapter",
     "extensions-builtin/sd-extension-chainner/nodes",
@@ -401,23 +402,24 @@ exclude = [
     "scripts/instantir/*",
     "scripts/softfill.py",
     "scripts/custom_code.py",
-    "pipelines/zetachroma/",
-    "pipelines/xomni/",
-    "pipelines/vibe/",
-    "pipelines/ultraflux/",
-    "pipelines/step1x/",
-    "pipelines/omnigen2/",
-    "pipelines/meissonic/",
-    "pipelines/model_stablecascade.py",
-    "pipelines/lumina_dimmo",
-    "pipelines/hidream",
-    "pipelines/f_lite",
+    "pipelines/anima",
+    "pipelines/boogu",
     "pipelines/bria",
     "pipelines/ernie",
+    "pipelines/f_lite",
     "pipelines/flex2",
-    "pipelines/anima",
+    "pipelines/hidream",
+    "pipelines/lumina_dimmo",
     "pipelines/meissonic",
-    "pipelines/boogu",
+    "pipelines/meissonic/",
+    "pipelines/model_stablecascade.py",
+    "pipelines/omnigen2/",
+    "pipelines/sefi/",
+    "pipelines/step1x/",
+    "pipelines/ultraflux/",
+    "pipelines/vibe/",
+    "pipelines/xomni/",
+    "pipelines/zetachroma/",
     "extensions-builtin/sd-extension-chainner/nodes",
 ]
 

scripts/color_grading.py

@@ -16,6 +16,8 @@ class ScriptPostprocessingColorGrading(scripts_postprocessing.ScriptPostprocessi
         grading_params = processing_grading.GradingParams(*args, **kwargs)
         if not processing_grading.is_active(grading_params):
             return
+        if pp.image is None:
+            return
         pp.image = processing_grading.grade_image(pp.image, grading_params)
         defaults = processing_grading.GradingParams()
         for f in fields(grading_params):

scripts/postprocessing_detailer.py

@@ -28,6 +28,10 @@ class ScriptPostprocessingDetailer(scripts_postprocessing.ScriptPostprocessing):
                 sampler='Default', prediction='default', shift=3.0, cfg_scale=6.0, options=None, seed=-1):
         if not enabled:
             return pp
+        if not shared.sd_loaded:
+            log.warning('Detailer postprocess: SD model not loaded')
+            pp.info["Detailer"] = "skipped (SD model not loaded)"
+            return pp
         if shared.sd_model is None or not hasattr(shared.sd_model, 'sd_checkpoint_info'):
             log.warning('Detailer postprocess: no base model selected')
             pp.info["Detailer"] = "skipped (no base model selected)"