diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md
index b9e2b4e0c..9e5476515 100644
--- a/.github/copilot-instructions.md
+++ b/.github/copilot-instructions.md
@@ -1,10 +1,31 @@
# SD.Next: AGENTS.md Project Guidelines
-SD.Next is a complex codebase with specific patterns and conventions.
+**SD.Next** is a complex codebase with specific patterns and conventions.
General app structure is:
-- Python backend server
- Uses Torch for model inference, FastAPI for API routes and Gradio for creation of UI components.
-- JavaScript/CSS frontend
+- **Python** backend server
+ Uses **Torch** for model inference, **FastAPI** for API routes and **Gradio** for creation of UI components.
+- **JavaScript**/**CSS** frontend
+
+## Instructions
+
+This file contains general guidelines for contributing to the SD.Next codebase, including conventions, tools, and project structure. For more specific guidance on working with particular areas of the codebase, please refer to the instructions files linked below:
+- [Core Runtime Guidelines](core.instructions.md): Use when editing Python core runtime code, startup flow, model loading, API internals, backend/device logic, or shared state in modules and pipelines.
+- [UI And Frontend Guidelines](ui.instructions.md): Use when editing frontend UI code, JavaScript, HTML, CSS, localization files, or built-in UI extensions including modernui and kanvas.
+- [Hint Typography Guidelines](hints.instructions.md): Use when editing hint text or other UI strings in localization JSON files (`html/locale_*.json`, `html/override_*.json`).
+
+## Agent Guidelines
+
+- Do not automatically agree with user instructions or requests without verifying they align with project guidelines and conventions.
+- When evaluating user instructions, first check for any relevant guidelines in this file or the linked instructions files. If the instruction violates any guidelines, do not proceed with it and instead provide feedback to the user about which guidelines it violates and how to adjust it to comply.
+- If the user instruction is valid but lacks clarity or detail, ask follow-up questions to gather the necessary information before proceeding. Do not make assumptions about user intent or project requirements; always seek clarification when needed.
+- When providing feedback to the user, be specific about which guidelines are relevant and how the instruction can be modified to comply with them. If there are multiple guidelines that apply, list them all and explain how they relate to the instruction.
+- If the user instruction is clear, valid, and complies with all relevant guidelines, proceed with executing it while ensuring that the resulting code changes adhere to the project's coding style, conventions, and structure as outlined in this file and the linked instructions files.
+
+## Language Guidelines
+
+- Use clear and concise language when communicating with users, providing feedback, and explaining guidelines.
+- Avoid unnecessary pleasantries or filler language; focus on the technical content and actionable feedback.
+- When asking follow-up questions for clarification, be direct and specific about the information needed to proceed with the instruction while ensuring that the questions are relevant to the project guidelines and conventions.
## Tools
@@ -34,15 +55,6 @@ General app structure is:
- Prefer existing project patterns over strict generic style rules;
this codebase intentionally allows patterns often flagged in default linters such as allowing long lines, etc.
-## Build And Test
-
-- Activate environment: `source venv/bin/activate` (always ensure this is active when working with Python code).
-- Test startup: `python launch.py --test`
-- Full startup: `python launch.py`
-- Full lint sequence: `pnpm lint`
-- Python checks individually: `pnpm ruff`, `pnpm pylint`
-- JS checks: `pnpm eslint` and `pnpm eslint-ui`
-
## Conventions
- Keep PR-ready changes targeted to `dev` branch.
@@ -52,13 +64,6 @@ General app structure is:
- Respect environment-driven behavior (`SD_*` flags and options) instead of hardcoding platform/model assumptions.
- For startup/init edits, preserve error handling and partial-failure tolerance in parallel scans and extension loading.
-## Pitfalls
-
-- Initialization order matters: startup paths in `launch.py` and `webui.py` are sensitive to import/load timing.
-- Shared mutable global state can create subtle regressions; prefer narrow, explicit changes.
-- Device/backend-specific code paths (**CUDA/ROCm/IPEX/DirectML/OpenVINO**) should not assume one platform.
-- Scripts and extension loading is dynamic; failures may appear only when specific extensions or models are present.
-
## File Creation
- Any temporary scripts or markdown reports must be stored in `tmp/` folder
@@ -73,6 +78,10 @@ Use these repo-local skills for recurring SD.Next model integration work:
File: `.github/skills/port-model/SKILL.md`
Use when adding a new model family, porting a standalone script into a Diffusers pipeline, or wiring an upstream Diffusers model into SD.Next.
+- `port-pipeline`
+ File: `.github/skills/port-pipeline/SKILL.md`
+ Use when porting a custom model pipeline implementation to a Diffusers pipeline class with behavior parity and no hard-coded device or attention assumptions.
+
- `debug-model`
File: `.github/skills/debug-model/SKILL.md`
Use when a new or existing SD.Next/Diffusers model integration fails during detection, loading, prompt encoding, sampling, or output handling.
diff --git a/.github/instructions/core.instructions.md b/.github/instructions/core.instructions.md
index 6d40c80bb..076bb5008 100644
--- a/.github/instructions/core.instructions.md
+++ b/.github/instructions/core.instructions.md
@@ -13,3 +13,19 @@ applyTo: "launch.py, webui.py, installer.py, modules/**/*.py, pipelines/**/*.py,
- Follow existing API/server patterns under `modules/api/` and reuse shared queue/state helpers rather than ad-hoc request handling.
- Reuse established model-loading and pipeline patterns (`modules/sd_*`, `pipelines/`) instead of creating parallel abstractions.
- For substantial Python changes, run at least relevant checks: `npm run ruff` and `npm run pylint` (or narrower equivalents when appropriate).
+
+## Build And Test
+
+- Activate environment: `source venv/bin/activate` (always ensure this is active when working with Python code).
+- Test startup: `python launch.py --test`
+- Full startup: `python launch.py`
+- Full lint sequence: `pnpm lint`
+- Python checks individually: `pnpm ruff`, `pnpm pylint`
+- JS checks: `pnpm eslint` and `pnpm eslint-ui`
+
+## Pitfalls
+
+- Initialization order matters: startup paths in `launch.py` and `webui.py` are sensitive to import/load timing.
+- Shared mutable global state can create subtle regressions; prefer narrow, explicit changes.
+- Device/backend-specific code paths (**CUDA/ROCm/IPEX/DirectML/OpenVINO**) should not assume one platform.
+- Scripts and extension loading is dynamic; failures may appear only when specific extensions or models are present.
diff --git a/.github/instructions/hints.instructions.md b/.github/instructions/hints.instructions.md
new file mode 100644
index 000000000..8243bdb53
--- /dev/null
+++ b/.github/instructions/hints.instructions.md
@@ -0,0 +1,50 @@
+---
+description: "Use when editing hint text or other UI strings in localization JSON files."
+name: "Hint Typography Guidelines"
+applyTo: "html/locale_*.json, html/override_*.json"
+---
+# Hint Typography Guidelines
+
+Hint strings render as HTML. Use this small set of inline tags to keep hints scannable:
+
+- `` for values: defaults, dropdown enums, specific numerics. Examples: `0.30 `, `Karras `, `v_prediction `, `UniPC `.
+- `... ` for cross-references to other UI controls by their exact visible label. Examples: `Denoising strength `, `Use init image `, `Images tab.
+- `` for proper nouns: model families, datasets, technique names. Examples: `SDXL `, `Flux `, `ControlNet `, `YOLO `.
+- `` for literals: paths, filename tokens, command-line snippets to type or use verbatim. Examples: `models/yolo`, `-seg`, `[PROMPT]`.
+
+## Cross-references
+
+- Use `... ` whenever a hint refers to another control by its exact visible label. This includes setting names, tab names, and named buttons.
+- Match the label exactly, including capitalization and spacing; readers look for the same string in the UI.
+- Do not use `` and `` separately for cross-references; always combine them.
+- Generic concept references (`the model`, `the prompt`, `the scheduler`) stay unstyled.
+
+## Tab naming
+
+- Refer to the unified generation tab as `Images ` (the ModernUI label). Do not write "Control tab"; that label only exists in legacy Standard UI.
+- "Control" remains valid as a setting value (`No: Control only `) or as part of a UI element name (`Control input ` pane), just not as a tab name.
+
+## Structure
+
+- ` ` for a line break within a paragraph.
+- ` ` for a paragraph break.
+- ` - key : description` for a keyed bullet list, used for short enumerations of dropdown values, modes, or numeric brackets. Each bullet's key is bolded; descriptions stay plain.
+- Do not use ``, ``, Markdown asterisks, or unicode bullets.
+
+## Common pitfalls
+
+- Do not bold ad-hoc emphasis; `` is reserved for values and, combined with ``, for cross-references.
+- Do not use `` for filenames, paths, or command tokens; those are literals and use ``.
+- Do not reword the inside of `` blocks; they are literal user-facing strings.
+- Stay ASCII; prefer semicolons or two sentences over em-dashes. The locale file convention is ASCII-only.
+
+## Translation propagation
+
+- `html/locale_en.json` is the source of truth. Other `html/locale_*.json` files are auto-generated by `cli/localize.js`; edit only the English file.
+- Per-locale corrections live in `html/override_{locale}.json`.
+
+## Validation
+
+- Validate JSON syntax with `jq empty html/locale_en.json`.
+- Lint with `pnpm eslint -- html/locale_en.json` (silent success).
+- See `wiki/Hints.md` for the wiki-facing version of these rules.
diff --git a/.github/skills/README.md b/.github/skills/README.md
index a1f4a3c70..5c08ec501 100644
--- a/.github/skills/README.md
+++ b/.github/skills/README.md
@@ -8,6 +8,10 @@ This folder contains repo-local Copilot skills for recurring SD.Next tasks.
File: `port-model/SKILL.md`
Use when adding or porting a model family into SD.Next and Diffusers.
+- `port-pipeline`
+ File: `port-pipeline/SKILL.md`
+ Use when porting a custom pipeline implementation into a Diffusers pipeline class while preserving behavior and avoiding hard-coded runtime assumptions.
+
- `debug-model`
File: `debug-model/SKILL.md`
Use when a new or existing SD.Next/Diffusers model integration fails during detect, load, prompt encode, sample, or output handling.
diff --git a/.github/skills/port-model/SKILL.md b/.github/skills/port-model/SKILL.md
index a9c223baf..fe1f0a2ac 100644
--- a/.github/skills/port-model/SKILL.md
+++ b/.github/skills/port-model/SKILL.md
@@ -63,6 +63,11 @@ Before implementing model-reference updates, explicitly ask the user which categ
Do not guess this category. Use the user answer to decide which reference JSON file(s) to update.
+## Mandatory Pipeline Question
+
+Before implementing a pipeline, explicitly ask the user if the model already has an upstream Diffusers pipeline that can be reused.
+If not, ask for URL or path to a reference implementation that can be structurally copied.
+
## Repo Files To Check
Start by reading the task description, then inspect the closest matching implementations.
@@ -130,6 +135,9 @@ Pipeline module responsibilities:
- Output dataclass
- Optional callback handling and output conversion
+If custom pipeline is provided by user, check it for accuracy and completness but do not assume it is perfect. Make necessary adjustments to fit SD.Next patterns and validate the result.
+Fix all relative imports to be absolute and compatible with SD.Next repo structure, make sure that all imports are resolvable and make sure it passes `ruff` checks.
+
### 3. Raw Checkpoint Or Single-File Weights
Use this path when the model source is not a normal Diffusers repository.
diff --git a/.github/skills/port-pipeline/SKILL.md b/.github/skills/port-pipeline/SKILL.md
new file mode 100644
index 000000000..2b18599c9
--- /dev/null
+++ b/.github/skills/port-pipeline/SKILL.md
@@ -0,0 +1,102 @@
+---
+name: port-pipeline
+description: "Port custom model pipeline implementations to Diffusers. Use when migrating custom or non-Diffusers pipeline code into SD.Next repo-local pipeline files such as pipelines/model_.py or pipelines//pipeline.py while preserving behavior, avoiding new dependencies, and keeping device/attention handling configurable."
+argument-hint: "Provide source pipeline path, target SD.Next destination path, and target pipeline class name"
+---
+
+# Port Custom Pipeline To Diffusers
+
+Port an existing custom model pipeline implementation into a Diffusers-compatible pipeline class with behavior parity and SD.Next-friendly conventions.
+This skill targets SD.Next repo-local pipeline ports only.
+
+## When To Use
+
+- A user has a custom pipeline implementation and wants it ported to Diffusers
+- Existing model code is runnable but not structured as a Diffusers pipeline
+- The destination is SD.Next pipeline code under `pipelines/model_*.py` or `pipelines//`
+- The task requires preserving generation behavior without introducing new dependencies
+- The task requires removing hard-coded runtime assumptions (device or attention backend)
+
+## Mandatory Clarification Gate
+
+Before implementation, confirm these required inputs with the user:
+
+1. Path to the source custom pipeline implementation
+2. Destination path in this SD.Next repository (typically under `pipelines/`)
+3. Target pipeline class name
+
+If any of the above are missing or ambiguous, stop and ask concise clarification questions before writing code.
+
+## Constraints
+
+- Do not add new dependencies
+- Do not hard-code device type (`cpu`, `cuda`, `mps`, etc.)
+- Do not hard-code attention type or backend assumptions
+- Preserve externally visible behavior of the source pipeline unless the user asks for intentional changes
+
+## Workflow
+
+1. Collect Inputs
+- Ask for source path, destination path, and target pipeline name.
+- Confirm destination is an SD.Next repo-local pipeline location, not an upstream Diffusers repository path.
+- Confirm runtime assumptions and expected task type (text-to-image, image-to-image, inpaint, etc.).
+
+2. Analyze Source Pipeline
+- Inspect model loading, prompt processing, denoising or sampling loop, scheduler interactions, and output post-processing.
+- Identify all components that must be ported: models, tokenizers or processors, schedulers, adapters, preprocessors, postprocessors, callbacks, and output dataclasses.
+- Note any hidden global state, side effects, or implicit defaults that must become explicit parameters.
+
+3. Map To Diffusers Interfaces
+- Choose the most appropriate Diffusers base class and output type.
+- Define `__init__`, module registration, `from_pretrained` and `__call__` signatures aligned with existing Diffusers patterns.
+- Keep parameter names and behavior as close as possible to upstream conventions.
+- Identify any custom classes needed beyond the pipeline itself: transformer blocks, attention processors, custom schedulers, or output types. Plan a separate module file for each.
+
+4. Implement Supporting Classes
+- If the pipeline requires custom model classes (e.g., a custom transformer block, attention module, or other model component), implement each in a **separate module** located in the **same directory** as the main pipeline file (e.g., `pipelines//transformer.py`, `pipelines//scheduler.py`).
+- If the pipeline requires a custom scheduler class, implement it in its own module (e.g., `pipelines//scheduler_.py`) following Diffusers scheduler conventions (`step`, `add_noise`, `scale_model_input`, etc.).
+- Each supporting class module must be self-contained: no circular imports, no hidden global state, and no hard-coded device or attention assumptions.
+- Import supporting classes into the main pipeline module from their respective sibling modules.
+
+5. Implement Pipeline Class
+- Create the destination pipeline classes at the user-provided path.
+- Port logic in small, testable sections: initialization, input validation, prompt encoding, latent preparation, denoising loop, decoding, and output packaging.
+- Replace hard-coded device and attention logic with runtime-configurable behavior.
+- Keep imports limited to existing project and Diffusers dependencies.
+
+6. Lint And Fix
+- Activate the project venv: `source venv/bin/activate`
+- Run `ruff` on all newly written files: `pnpm ruff` (or `ruff check --fix` for targeted runs).
+- Run `pylint` on all newly written files: `pnpm pylint` (or `pylint ` for targeted runs).
+- Fix every reported error or warning that is not explicitly marked with a `TODO` suppression comment in the source.
+- Re-run both linters after fixes to confirm a clean result before proceeding.
+
+7. Validate Behavior Parity
+- Compare source and ported implementations for input-output shape handling, dtype flow, scheduler step ordering, and guidance behavior.
+- Run focused checks or smoke tests if available in the workspace.
+- Call out any known differences that were required for Diffusers compatibility.
+
+8. Report Results
+- Summarize what was ported and where.
+- List any unresolved assumptions, risks, or TODOs.
+- Provide minimal follow-up steps for integration and testing.
+
+## Review Checklist
+
+- Required inputs were collected before edits
+- No new dependency was introduced
+- No hard-coded device or attention backend remains
+- Core components from source pipeline were fully mapped
+- Pipeline class is in requested destination with requested name
+- Each custom supporting class (transformer, scheduler, etc.) is in its own sibling module
+- Supporting modules have no circular imports or hidden global state
+- `ruff` and `pylint` both pass cleanly on all newly written files (venv activated)
+- Main inference path behavior matches the source implementation
+
+## Output Expectations
+
+Final response should include:
+- Source path, destination path, and final pipeline class name
+- Brief parity summary of key components ported
+- Validation performed and any gaps
+- Explicit note of any assumptions requiring user confirmation
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29553c783..b3f98aa8e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,122 @@
# Change Log for SD.Next
+## Update for 2026-05-13
+
+### Highlights for 2026-05-13
+
+Just two weeks since last release, but we have a lot of new models and features to cover!
+
+*What's New?*
+- Image editing models now can work with multiple image inputs!
+- Six new models: *HiDream-O1 Image*, *JoyAI Image Edit*, *Step1X-Edit*, *VIBE Image Edit* and *UltraFlux*
+- Enhanced capabilities for *Anima*, *Ernie-Image*, *LTX*, *Flux.2* and *Chroma* models
+- Enhanced *LoRA* capabilities in many models
+- UI improvements across the board: *Main panels*, *Gallery*, *Kanvas*, *Networks*, and more...
+
+For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md)
+
+[ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic)
+
+### Details for 2026-05-13
+
+- **Models**
+ - [HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) pixel-level unified transformer model support
+ HiDream-O1 is based on a single custom *Qwen3-VL* 8.8B 35GB component
+ includes both **HiDream-O1-Image** *(base)* and **HiDream-O1-Image-Dev** *(distilled*)* variants
+ includes *sdnq-svd-dynamic-int8* pre-quantized variants for both base and dev models
+ includes *T2I* and *I2I edit* capabilities and resolutions up to 2048px
+ *note*: use steps:50 for base and steps:28 for dev variants
+ - [JoyAI Image Edit](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Diffusers) image-editing model support
+ includes multimodal conditioning using *Qwen3-VL* with a dedicated *JoyImageEdit* diffusion transformer
+ *note* this is a large model at 50GB so use of aggressive quantization is recommended
+ - [StepFun Step1X-Edit v1.1](https://huggingface.co/stepfun-ai/Step1X-Edit-v1p1-diffusers) image-editing model support
+ step1x is a large dedicated image edit model combining qwen-2.5 8B encoder with custom 12.4B transformer
+ - [VIBE Image Edit](https://huggingface.co/iitolstykh/VIBE-Image-Edit) image-editing model support
+ built on Sana1.5-1.6B diffusion backbone with Qwen3-VL-2B multimodal conditioning
+ primarily image-editing model, but supports t2i as well, uses multi-scale resolution binning up to 2048px
+ - [AlphaVLLM Lumina-DiMOO](https://huggingface.co/Alpha-VLLM/Lumina-DiMOO) unified multimodal diffusion model
+ includes *T2I*, *I2I edit*, and *MMU* capabilities in a single pipeline
+ *note* model also supports special prompts: *dense, canny_pred, control, subject, edit, ref_transfer, multi_view*
+ *note* as with most multimodal/unified models, it needs higher step count (recommended is 64 steps) and uses quite a lot of VRAM, so use with caution!
+ - [Owen777 UltraFlux-v1](https://huggingface.co/Owen777/UltraFlux-v1) native 4K text-to-image model based on *FLUX.1-dev*
+ *note*: UltraFlux is capable of rendering images up to 4K resolution, but it doesnt mean it will do that on any hardware - it will depend on your VRAM!
+ - [Anima Preview-v3](https://huggingface.co/circlestone-labs/Anima)
+ add *turbo* variant with [turbo-LoRA](https://civitai.com/models/2560840/anima-turbo-lora) pre-merged
+ add *sdnq-svd-dynamic-int8* pre-quantized variant
+- **Features**
+ - **Multi-image** workflows!
+ for models that support multiple images as inputs, you can now add multiple stages in Kanvas
+ prompts like "*place character from first image, add background from second image, render in style from third image*" are now possible
+ - option *inputs -> skip processing* to force images to passed to model as-is without any pre-processing
+ examples of models that support multi-inputs: *qwen-image-edit, flux.2, google-gemini*
+ - [SD Ultimate Upscale](https://github.com/Coyote-A/ultimate-upscale-for-automatic1111)
+ still a popular method for upscaling, but has not been updated nor maintained for a while
+ so now its modernized and fully integrated as a built-in script!
+ - **LTX** support for *audio* generation
+ - **Anima** support for *img2img* and *inpaint* workflows
+ - **Ernie-Image** add native *LoRA* support, *img2img* and *inpaint* workflows
+ - **Chroma** add native *LoRA* support
+ - **Flux.2** add native *LoRA* support
+ - **Prompt enhance** add info to image metadata
+ - custom **VAE** loader for all pipelines
+ *note*: vae still needs to be compatible with the model
+ - **CivitAI** downloaded thumbnails now include metadata
+ - **Installer** support for `git+http` style references
+- **UI**
+ - **Networks** using networks to load model or auto-download a reference model will now be reflected in the UI
+ - ability to manually reorient *input/output* panels
+ - all ui panels can be *minimized/maximized* by clicking on their header
+ state is preserved across sessions and can be used to hide rarely used panels and declutter the workspace
+ - **Kanvas** re-order stages by clicking on active stage
+ order of stages determines order of images passed to model
+ - **Kanvas** *magic-wand* tool now works on mask layer and auto-creates mask based on perceptual tolerance
+ - **Gallery** add thumbnail size slider
+ - **Gallery** add quick info/download/delete buttons on thumbnail hover
+ - **Models** sortable columns, ability to remove a model
+ applies to models as well as huggingface cache entries
+ - **Server Info** add button *copy-to-clipboard*
+ useful for sharing your system info when asking for help in discord or github
+- **Control**
+ - remove buttons: *input/control/process*
+ - move params *control input type* to control menu section
+ - remove "processed preview" from ui
+ preprocessor output can still be generated by clicking preview button in in control unit and it will render into normal output area
+- **Internal**
+ - `offload` auto-reapply hook on error
+ - refactor `pip` installer, thanks @awsr
+ - remove obsolete `lora` stepwise and functional code, thanks @awsr
+ - interrupt model loading between components
+ - patch `rich` for cleaner exception logging
+ - lint `ruff` strict and reduce exceptions
+ - lint `pylint` improvements
+ - lint `ty` readiness
+- **Fixes**
+ - add missing `jquery` and `sparkline` js scripts
+ - save handle already decoded images
+ - `ernie-image` preview
+ - `lora` false deactivate
+ - `kandinsky-5` t2i/i2i workflows
+ - progress do not timeout when paused
+ - faster server shutdown/restart, thanks @awsr
+ - `openvino` force offload none
+ - `lut` file handling
+ - warn on pipeline ignoring `cfg`
+ - detailer `segmentation`, thanks @awsr
+ - `ipex` invalid device type
+ - cache network thumbnails
+ - `scripts` corrupting control ui state
+ - avoid `callback` duplicate registrations
+ - pipeline task change causing loss of info on loaded `lora`
+ - `detailer` handle `lora` internally
+ - vae preview flashes previous image
+ - `torch.compile` improvements
+ - `gradio` preprocess exception handling
+ - `ipadapters` with offloading
+ - `kanvas` outpaint
+ - `network` preview handle invalid image
+ - `schedulers` improve *set_timesteps* handling
+ - `schedulers` improve *scale_noise* handling
+
## Update for 2026-04-28
### Highlights for 2026-04-28
diff --git a/README.md b/README.md
index 9f1a72a41..baa72b88e 100644
--- a/README.md
+++ b/README.md
@@ -57,7 +57,7 @@ SD.Next is feature-rich with a focus on performance, flexibility, and user exper
SD.Next includes many features not found in other WebUIs, such as:
- **SDNQ**: State-of-the-Art quantization engine
- Use pre-quantized or run with quantizaion on-the-fly for up to 4x VRAM reduction with no or minimal quality and performance impact
+ Use pre-quantized or run with quantization on-the-fly for up to 4x VRAM reduction with no or minimal quality and performance impact
- **Balanced Offload**: Dynamically balance CPU and GPU memory to run larger models on limited hardware
- **Captioning** with 150+ **OpenCLiP** models, **Tagger** with **WaifuDiffusion** and **DeepDanbooru** models, and 25+ built-in **VLMs**
- **Image Processing** with full image correction color-grading suite of tools
diff --git a/TODO.md b/TODO.md
index 3e58efbe8..dbc4230e5 100644
--- a/TODO.md
+++ b/TODO.md
@@ -4,13 +4,13 @@
### Assigned
-- Gallery: thumb-size, quick delete/download/info @vladmandic
- Chat-based interface, @vladmandic
-- Multi-image inputs, @vladmandic
- Control tab verify overrides handling, @vladmandic
- Reimplement `llama` remover for Kanvas, @vladmandic
-- Integrate [Depth3D](https://github.com/vladmandic/sd-extension-depth3d), @vladmandic
+- Implement [pruna](https://github.com/PrunaAI/pruna), @vladmandic
+- Change params to default, @vladmandic
+- Detailer postprocessing, @CalamitousFelicitousness
- Cloud providers, @CalamitousFelicitousness
- Video processing add full API support, @CalamitousFelicitousness
@@ -20,7 +20,6 @@
- `RIFE` in processing
- `SeedVR2` in processing
- Video model loader: Add video models to Reference
-- REMBG add
- UI Lite vs Expert mode
- TensorRT acceleration
- Auto handle scheduler `prediction_type`
@@ -56,21 +55,13 @@ TODO: Investigate which models are diffusers-compatible and prioritize!
### Image
- [JoyAI-Image-Edit](https://github.com/huggingface/diffusers/pull/13444) (pr in-progress)
-- [Lumina-DiMOO](https://github.com/huggingface/diffusers/pull/12468) (pr stalled)
-- [Step1X-Edit](https://github.com/huggingface/diffusers/pull/12249) (pr stalled)
-- [VIBE Image Edit](https://huggingface.co/iitolstykh/VIBE-Image-Edit) (diffusers-compatible)
- [nVidia Cosmos-Predict-2.5](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B) (in diffusers)
- [nVidia Cosmos-Transfer-2.5](https://huggingface.co/nvidia/Cosmos-Transfer2.5-2B) (in diffusers)
-- [UltraFlux](https://huggingface.co/Owen777/UltraFlux-v1) (diffusers-compatible)
- [Tencent HY-WU](https://huggingface.co/tencent/HY-WU) (transformers-compatible)
-- [Mugen](https://huggingface.co/CabalResearch/Mugen) (sdxl with flux vae experiment, not clean)
-- [Liquid](https://github.com/FoundationVision/Liquid) (autoregressive, not clean)
### Video
- [HY-OmniWeaving](https://huggingface.co/tencent/HY-OmniWeaving)
-- [LTX-Condition](https://huggingface.co/Lightricks/LTX-2)
-- [LTX-Distilled](https://huggingface.co/Lightricks/LTX-2)
- [OpenMOSS MOVA](https://huggingface.co/OpenMOSS-Team/MOVA-720p)
- [Wan2.2-Animate](https://huggingface.co/Wan-AI/Wan2.2-Animate-14B)
- [Wan2.1-T2V-14B-CausVid](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-CausVid)
@@ -84,12 +75,11 @@ TODO: Investigate which models are diffusers-compatible and prioritize!
- [Sana I2V](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p_diffusers)
- [Wan-2.2 S2V](https://huggingface.co/Wan-AI/Wan2.2-S2V-14B)
- [Meituan LongCat-Video](https://huggingface.co/meituan-longcat/LongCat-Video)
-- [LTXVideo LongMulti](https://huggingface.co/Lightricks/LTX-Video-0.9.8-13B-distilled)
- [Phantom HuMo](https://github.com/Phantom-video/Phantom)
- [CausVid-Plus](https://github.com/goatWu/CausVid-Plus/)
- [LivePortrait](https://github.com/KwaiVGI/LivePortrait)
-- [Magi (SandAI)](https://github.com/SandAI-org/MAGI-1)
-- [Ming (inclusionAI)](https://github.com/inclusionAI/Ming)
+- [SandAI Magi](https://github.com/SandAI-org/MAGI-1)
+- [inclusionAI Ming](https://github.com/inclusionAI/Ming)
- [HummingbirdXT](https://huggingface.co/amd/HummingbirdXT)
- [DiffusionForcing](https://github.com/kwsong0113/diffusion-forcing-transformer)
- [ByteDance Lynx](https://github.com/bytedance/lynx)
@@ -155,20 +145,30 @@ TODO: Investigate which models are diffusers-compatible and prioritize!
- Background removal model trained on Bria FIBO dataset
- Created: 2025-08 | Updated: 2025-09 | Stars: N/A (private model)
+### Rejected
+
+- [Mugen](https://huggingface.co/CabalResearch/Mugen) (sdxl with flux vae experiment, not clean)
+- [Liquid](https://github.com/FoundationVision/Liquid) (autoregressive, not clean)
+
## Code TODO
> npm run todo
```code
-installer.py:642:15: W0511: TODO rocm: switch to pytorch source when it becomes available (fixme)
-modules/transformer_cache.py:29:61: W0511: TODO fc: autodetect tensor format based on model (fixme)
-modules/transformer_cache.py:30:50: W0511: TODO fc: autodetect distilled based on model (fixme)
-modules/processing_class.py:404:32: W0511: TODO processing: remove duplicate mask params (fixme)
-modules/sd_samplers_diffusers.py:355:31: W0511: TODO enso-required (fixme)
-modules/sd_models.py:1356:5: W0511: TODO model load: implement model in-memory caching (fixme)
-modules/ui_models_load.py:257:5: W0511: TODO loader: load receipe (fixme)
-modules/ui_models_load.py:264:5: W0511: TODO loader: save receipe (fixme)
-modules/sd_hijack_hypertile.py:123:17: W0511: TODO hypertile: vae breaks when using non-standard sizes (fixme)
-modules/sd_unet.py:77:39: W0511: TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage (fixme)
-modules/modular_guiders.py:66:51: W0511: TODO: guiders (fixme)
+installer.py:TODO rocm: switch to pytorch source when it becomes available
+modules/control/run.py:TODO modernui: monkey-patch for missing tabs.select event
+modules/history.py:TODO: apply metadata, preview, load/save
+modules/image/resize.py:TODO resize image: enable full VAE mode for resize-latent
+modules/lora/lora_load.py:TODO lora: add t5 key support for sd35/f1
+modules/masking.py:TODO: additional masking algorithms
+modules/modular_guiders.py:TODO: guiders
+modules/processing_class.py:TODO processing: remove duplicate mask params
+modules/sd_hijack_hypertile.py:TODO hypertile: vae breaks when using non-standard sizes
+modules/sd_models.py:TODO model load: implement model in-memory caching
+modules/sd_samplers_diffusers.py:TODO enso-required
+modules/sd_unet.py:TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage
+modules/transformer_cache.py:TODO fc: autodetect distilled based on model
+modules/transformer_cache.py:TODO fc: autodetect tensor format based on model
+modules/ui_models_load.py:TODO loader: load receipe
+modules/ui_models_load.py:TODO loader: save receipe
```
diff --git a/data/reference-distilled.json b/data/reference-distilled.json
index ad7b07ddb..4588ae5c1 100644
--- a/data/reference-distilled.json
+++ b/data/reference-distilled.json
@@ -36,6 +36,16 @@
"skip": true,
"extras": "sampler: Default, cfg_scale: 4.5"
},
+ "HiDream-O1 Image Dev": {
+ "path": "HiDream-ai/HiDream-O1-Image-Dev",
+ "preview": "HiDream-ai--HiDream-O1-Image-Dev.jpg",
+ "desc": "HiDream-O1-Image-Dev is the distilled 8B HiDream-O1 variant tuned for 28-step fast generation using flash flow scheduling.",
+ "skip": true,
+ "extras": "sampler: Flash, steps: 28, cfg_scale: 0.0",
+ "size": 35.2,
+ "tags": "distilled",
+ "date": "2026 May"
+ },
"Qwen-Image-Lightning": {
"path": "vladmandic/Qwen-Lightning",
"preview": "vladmandic--Qwen-Lightning.jpg",
diff --git a/data/reference-quant.json b/data/reference-quant.json
index edeb0dbf2..67254baf4 100644
--- a/data/reference-quant.json
+++ b/data/reference-quant.json
@@ -242,7 +242,7 @@
"desc": "ERNIE-Image is a text-to-image diffusion transformer model that combines a Mistral3 text encoder with a FlowMatch transformer and Flux2-style VAE for 1024px image generation.",
"skip": true,
"extras": "sampler: Default, cfg_scale: 4.0, steps: 50",
- "size": 23.93,
+ "size": 7.52,
"date": "2026 April"
},
"Baidu ERNIE-Image-Turbo sdnq-dynamic-int4": {
@@ -251,8 +251,44 @@
"desc": "ERNIE-Image-Turbo is a distilled ERNIE-Image variant optimized for fast generation with fewer denoising steps.",
"skip": true,
"extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
- "size": 23.37,
+ "size": 7.52,
"tags": "distilled",
"date": "2026 April"
+ },
+ "Anima Preview 3 sdnq-dynamic-int8": {
+ "path": "vladmandic/Anima-Preview-3-diffusers-SDNQ-8bit-dynamic",
+ "preview": "CalamitousFelicitousness--Anima-Preview-3-sdnext-diffusers.jpg",
+ "desc": "Anima Preview V3 with extended 1024-resolution training and expanded dataset coverage for less common artists. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
+ "tags": "community",
+ "date": "2026 April",
+ "size": 3.19,
+ "skip": true
+ },
+ "Anima Preview 3 Turbo sdnq-dynamic-int8": {
+ "path": "vladmandic/Anima-Preview-3-turbo-diffusers-SDNQ-8bit-dynamic",
+ "preview": "CalamitousFelicitousness--Anima-Preview-3-sdnext-diffusers.jpg",
+ "desc": "Anima Preview V3 with extended 1024-resolution training and expanded dataset coverage for less common artists. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
+ "tags": "community",
+ "date": "2026 April",
+ "size": 3.19,
+ "skip": true
+ },
+ "HiDream-O1 Image sdnq-dynamic-int8": {
+ "path": "vladmandic/HiDream-O1-Image-SDNQ-8bit-dynamic",
+ "desc": "HiDream-O1-Image is an 8B pixel-level unified transformer model for text-to-image generation, instruction editing, and multi-reference personalization up to 2048x2048.",
+ "preview": "HiDream-ai--HiDream-O1-Image.jpg",
+ "skip": true,
+ "extras": "sampler: Default",
+ "size": 10.34,
+ "date": "2026 May"
+ },
+ "HiDream-O1 Image Dev sdnq-dynamic-int8": {
+ "path": "vladmandic/HiDream-O1-Image-Dev-SDNQ-8bit-dynamic",
+ "desc": "HiDream-O1-Image is an 8B pixel-level unified transformer model for text-to-image generation, instruction editing, and multi-reference personalization up to 2048x2048.",
+ "preview": "HiDream-ai--HiDream-O1-Image.jpg",
+ "skip": true,
+ "extras": "sampler: Default",
+ "size": 10.34,
+ "date": "2026 May"
}
}
diff --git a/data/reference.json b/data/reference.json
index 8fe66e55d..0a00af7cf 100644
--- a/data/reference.json
+++ b/data/reference.json
@@ -143,6 +143,16 @@
"date": "2025 January"
},
+ "Owen777 UltraFlux-v1": {
+ "path": "Owen777/UltraFlux-v1",
+ "preview": "Owen777--UltraFlux-v1.jpg",
+ "desc": "UltraFlux-v1 is a FLUX.1-dev based text-to-image model optimized for native 4K and multi-aspect-ratio generation with improved composition consistency.",
+ "skip": true,
+ "extras": "sampler: Default, cfg_scale: 4.0, steps: 50",
+ "size": 33.0,
+ "date": "2025 November"
+ },
+
"Z-Image": {
"path": "Tongyi-MAI/Z-Image",
"preview": "Tongyi-MAI--Z-Image.jpg",
@@ -622,6 +632,15 @@
"size": 20.75,
"date": "2025 January"
},
+ "AlphaVLLM Lumina DiMOO": {
+ "path": "Alpha-VLLM/Lumina-DiMOO",
+ "desc": "Lumina-DiMOO is an omni diffusion large language model for multimodal generation and understanding with text-to-image, image editing, and multimodal understanding capabilities.",
+ "preview": "Alpha-VLLM--Lumina-DiMOO.jpg",
+ "skip": true,
+ "extras": "sampler: Default",
+ "size": 0,
+ "date": "2025 September"
+ },
"HiDream-I1 Fast": {
"path": "HiDream-ai/HiDream-I1-Fast",
@@ -650,6 +669,15 @@
"size": 58.4,
"date": "2025 April"
},
+ "HiDream-O1 Image": {
+ "path": "HiDream-ai/HiDream-O1-Image",
+ "desc": "HiDream-O1-Image is an 8B pixel-level unified transformer model for text-to-image generation, instruction editing, and multi-reference personalization up to 2048x2048.",
+ "preview": "HiDream-ai--HiDream-O1-Image.jpg",
+ "skip": true,
+ "extras": "sampler: Default",
+ "size": 35.2,
+ "date": "2026 May"
+ },
"HiDream-E1 Full": {
"path": "HiDream-ai/HiDream-E1-Full",
"desc": "HiDream-E1 is an image editing model built on HiDream-I1.",
@@ -924,6 +952,36 @@
"extras": "sampler: Default, cfg_scale: 3.5",
"size": 16.2,
"date": "2025 December"
+ },
+
+ "StepFun Step1X-Edit v1.1": {
+ "path": "stepfun-ai/Step1X-Edit-v1p1-diffusers",
+ "preview": "stepfun-ai--Step1X-Edit-v1p1-diffusers.jpg",
+ "desc": "Multimodal image editing model using Step1X transformer architecture with Qwen2.5-VL text encoding, trained with Flow Matching scheduler for high-quality in-context image edits and refinements.",
+ "skip": true,
+ "extras": "sampler: Default",
+ "size": 24.85,
+ "date": "2025 September"
+ },
+
+ "VIBE Image Edit": {
+ "path": "vladmandic/VIBE-Image-Edit",
+ "preview": "vladmandic--VIBE-Image-Edit.jpg",
+ "desc": "VIBE is an open-source text-guided image editing model combining Sana1.5-1.6B diffusion backbone with Qwen3-VL multimodal conditioning for fast, instruction-based edits.",
+ "skip": true,
+ "extras": "sampler: Default, cfg_scale: 4.5, image_guidance_scale: 1.2, steps: 20",
+ "size": 9.27,
+ "date": "2025 December"
+ },
+
+ "JoyAI Image Edit": {
+ "path": "jdopensource/JoyAI-Image-Edit-Diffusers",
+ "preview": "jdopensource--JoyAI-Image-Edit-Diffusers.jpg",
+ "desc": "JoyAI Image Edit is a Diffusers-native image editing model that combines a JoyImageEdit transformer with Qwen3-VL multimodal conditioning for instruction-guided edits.",
+ "skip": true,
+ "size": 50.31,
+ "extras": "sampler: Default",
+ "date": "2026 April"
}
}
diff --git a/eslint.config.mjs b/eslint.config.mjs
index 1b6120ae1..403755e76 100644
--- a/eslint.config.mjs
+++ b/eslint.config.mjs
@@ -102,6 +102,7 @@ const jsConfig = defineConfig([
idbPut: 'readonly',
idbDel: 'readonly',
idbAdd: 'readonly',
+ initTableSorter: 'readonly',
idbCount: 'readonly',
idbFolderCleanup: 'readonly',
idbClearAll: 'readonly',
diff --git a/extensions-builtin/sd-extension-chainner b/extensions-builtin/sd-extension-chainner
index d4eab2166..1ac7d87cf 160000
--- a/extensions-builtin/sd-extension-chainner
+++ b/extensions-builtin/sd-extension-chainner
@@ -1 +1 @@
-Subproject commit d4eab2166e4d9b52e42924cc942198f9e22eb916
+Subproject commit 1ac7d87cf70f51d958953133fb457d7084baf701
diff --git a/extensions-builtin/sdnext-kanvas b/extensions-builtin/sdnext-kanvas
index 9a76f2093..8aa27c614 160000
--- a/extensions-builtin/sdnext-kanvas
+++ b/extensions-builtin/sdnext-kanvas
@@ -1 +1 @@
-Subproject commit 9a76f209312dc88b267af126c60fc8b368678131
+Subproject commit 8aa27c6143f9cae3d3914dd9203918da555c5ca9
diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index f81771717..c5382ac49 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit f81771717424299e99542b2ceb88e3cca9f7f16c
+Subproject commit c5382ac496bf71dbd6a1100d35f4c3a204d0fd66
diff --git a/html/locale_ar.json b/html/locale_ar.json
index 596040792..b1c1d3cb5 100644
--- a/html/locale_ar.json
+++ b/html/locale_ar.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "وسائط الإدخال",
"reload": "",
"hint": "إضافة صورة إدخال لاستخدامها في معالجة التحويل من صورة إلى صورة، أو التلوين (Inpaint)، أو التحكم"
diff --git a/html/locale_bn.json b/html/locale_bn.json
index f7fe99ac3..252670494 100644
--- a/html/locale_bn.json
+++ b/html/locale_bn.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "ইনপুট মিডিয়া",
"reload": "",
"hint": "ইমেজ-টু-ইমেজ, ইনপেইন্ট বা কন্ট্রোল প্রসেসিংয়ের জন্য ইনপুট ছবি যোগ করুন"
diff --git a/html/locale_de.json b/html/locale_de.json
index 1bdd46a1a..9fa94e636 100644
--- a/html/locale_de.json
+++ b/html/locale_de.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "Eingabemedien",
"reload": "",
"hint": "Eingabebild hinzufügen, das für Image-to-Image-, Inpaint- oder Control-Verarbeitung verwendet werden soll"
diff --git a/html/locale_en.json b/html/locale_en.json
index db1f36ed2..e31678ded 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -24,7 +24,7 @@
{"id":"xyz_grid_x_list","label":"⊜","localized":"","hint":"Fill","ui":"script_xyz_grid_script"},
{"id":"txt2img_caption_output","label":"","localized":"","hint":"Caption image","ui":"txt2img"},
{"id":"txt2img_image_fit","label":"⁜","localized":"","hint":"Cycle image fit method","ui":"txt2img"},
- {"id":"","label":"➠ Control","localized":"","hint":"Transfer image to control interface. Right-click this button to transfer only the prompt or all generation parameters to the Images tab without sending the image itself.","ui":"txt2img"},
+ {"id":"","label":"➠ Control","localized":"","hint":"Transfer image to the Images tab. Right-click this button to transfer only the prompt or all generation parameters without sending the image itself.","ui":"txt2img"},
{"id":"","label":"➠ Text","localized":"","hint":"Transfer image to text interface","ui":"txt2img"},
{"id":"","label":"➠ Image","localized":"","hint":"Transfer image to image interface","ui":"txt2img"},
{"id":"","label":"➠ Process","localized":"","hint":"Transfer image to process interface","ui":"txt2img"},
@@ -72,16 +72,16 @@
{"id":"","label":"Advanced Options","localized":"","hint":"","ui":"settings_sd"},
{"id":"","label":"Appearance","localized":"","hint":"","ui":"settings_ui"},
{"id":"","label":"Answer","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"Adjust start","localized":"","hint":"Starting step when sigma adjust occurs","ui":"txt2img"},
- {"id":"","label":"Adjust end","localized":"","hint":"Ending step when sigma adjust occurs","ui":"txt2img"},
+ {"id":"","label":"Adjust start","localized":"","hint":"Lower bound of the denoising window where Sigma adjust is active, as a fraction of the noise schedule (1.0 = pure noise, 0.0 = clean image). The adjustment stops once denoising progresses past this point, so higher values end the effect earlier. Default 0.2 leaves the final ~20% of the schedule unmodified.","ui":"txt2img"},
+ {"id":"","label":"Adjust end","localized":"","hint":"Upper bound of the denoising window where Sigma adjust is active, as a fraction of the noise schedule (1.0 = pure noise, 0.0 = clean image). The adjustment only begins once denoising has progressed past this point, so lower values delay the effect further into the run. Default 0.8 leaves the first ~20% of the schedule unmodified.","ui":"txt2img"},
{"id":"","label":"Autocomplete","localized":"","hint":"Enable or disable Tag Autocomplete. Choose which dictionaries are used for prompt autocompletion in Extras","ui":"control"},
{"id":"","label":"AutoGuidance dropout","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"AutoGuidance layers","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"AutoGuidance config","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"APG momentum","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"APG rescale","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"Attention guidance","localized":"","hint":"CFG scale used for with PAG: Perturbed-Attention Guidance","ui":"txt2img"},
- {"id":"","label":"Adaptive scaling","localized":"","hint":"Adaptive modifier for attention guidance scale","ui":"txt2img"},
+ {"id":"","label":"Attention guidance","localized":"","hint":"Dual-purpose slider that activates one of two guidance mechanisms depending on the loaded model. - SD 1.5 and SDXL : enables Perturbed Attention Guidance (PAG). Sdnext silently swaps the pipeline to a PAG-aware variant and steers generation away from a self-attention-perturbed prediction, improving structure and detail. Used in addition to the regular Guidance scale . - Flux, QwenImage, HiDream, Hunyuan Video, Sana, and other flow-matching models : routes to true_cfg_scale, enabling classifier-free guidance with negative prompts on models that don't natively use CFG. On other models the slider has no effect. Set to 0 to disable. Disabled by default.","ui":"txt2img"},
+ {"id":"","label":"Adaptive scaling","localized":"","hint":"Decay rate for the Perturbed Attention Guidance (PAG) component of Attention guidance . Higher values cause PAG strength to decay faster across the denoising steps. Only takes effect on SD 1.5 and SDXL when Attention guidance is non-zero (the only path that actually enables PAG). Has no effect on Flux , QwenImage , HiDream , or other flow-matching models that route Attention guidance to true_cfg_scale instead. Default 0.5 applies moderate decay. Set to 0 to keep PAG at full strength for the entire process.","ui":"txt2img"},
{"id":"","label":"Apply to hires","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Active IP adapters","localized":"","hint":"Number of active IP adapter","ui":"txt2img"},
{"id":"","label":"Adapter","localized":"","hint":"IP adapter model","ui":"txt2img"},
@@ -104,9 +104,9 @@
{"id":"","label":"ACI: Mask blur","localized":"","hint":"Adjust blur to apply a smooth transition between image and inpainted area. (Recommended value = 0 for sharpness)","ui":"script_automatic_color_inpaint"},
{"id":"","label":"Adaptive restore","localized":"","hint":"","ui":"script_instantir"},
{"id":"","label":"Apply noise","localized":"","hint":"","ui":"script_softfill"},
- {"id":"","label":"Auto min score","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Auto-segment","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Auto-mask","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Auto min score","localized":"","hint":"Minimum stability score for masks produced by Auto-segment (Facebook SAM and SlimSAM models). Higher values keep only the most confident masks; lower values include more candidates including noisier ones. Has no effect on Rembg models or on the Auto-mask threshold/edge methods. Default 0.8.","ui":"control"},
+ {"id":"","label":"Auto-segment","localized":"","hint":"Automatic foreground segmentation model. Runs on the input image to generate a mask without manual painting.None : no auto-segmentation; the manually painted mask is used instead.Facebook SAM ViT (Base/Large/Huge) : Meta's Segment Anything Model. Quality scales with size, Huge is the most accurate but slowest and largest in VRAM.SlimSAM Uniform / Uniform Tiny : pruned, faster SAM variants with a good speed/quality tradeoff for repetitive workflows.Rembg BEN2 / Silueta / U2Net / U2Net human / ISNet general / ISNet anime : lightweight background-removal models. Pick by content: U2Net human or BEN2 for people, ISNet anime for illustrations, U2Net or Silueta for general subjects. Models are downloaded on first use. Default None.","ui":"control"},
+ {"id":"","label":"Auto-mask","localized":"","hint":"Automatic mask generation from the input image using simple computer-vision methods (no neural model). Runs only when no manual mask is painted and Auto-segment is set to None.None : disabled.Threshold : Otsu binary threshold; everything brighter than the auto-computed threshold becomes the mask. Works for high-contrast subjects on plain backgrounds.Edge : detects contours and keeps the largest ones; useful for masking distinct objects with clear outlines.Grayscale : uses the image's luminance as the mask intensity, producing a soft, gradient-style mask. Default None.","ui":"control"},
{"id":"","label":"Active","localized":"","hint":"","ui":"control"},
{"id":"","label":"Attention","localized":"","hint":"","ui":"control"},
{"id":"","label":"Adain","localized":"","hint":"","ui":"control"},
@@ -163,13 +163,13 @@
{"id":"","label":"BitsAndBytes","localized":"","hint":"","ui":"settings_quantization"},
{"id":"","label":"Batch count","localized":"","hint":"How many batches of images to create (has no impact on generation performance or VRAM usage)","ui":"txt2img"},
{"id":"","label":"Batch size","localized":"","hint":"How many image to create in a single batch (increases generation performance at cost of higher VRAM usage)","ui":"txt2img"},
- {"id":"","label":"Beta schedule","localized":"","hint":"Defines how beta (noise strength per step) grows. Options: - default: the model default - linear: evenly decays noise per step - scaled: squared version of linear, used only by Stable Diffusion - cosine: smoother decay, often better results with fewer steps - sigmoid: sharp transition, experimental","ui":"txt2img"},
+ {"id":"","label":"Beta schedule","localized":"","hint":"Defines how beta (noise strength per step) grows. Options: - default : the model default - linear : evenly decays noise per step - scaled : squared version of linear, used only by Stable Diffusion - cosine : smoother decay, often better results with fewer steps - sigmoid : sharp transition, experimental","ui":"txt2img"},
{"id":"","label":"Base shift","localized":"","hint":"Minimum shift value for low resolutions when using dynamic shifting.","ui":"txt2img"},
{"id":"","label":"Brightness","localized":"","hint":"Adjusts overall image brightness. Positive values lighten the image, negative values darken it. Applied uniformly across all pixels in linear space.","ui":"txt2img"},
{"id":"","label":"Block","localized":"","hint":"","ui":"script_kohya_hires_fix"},
{"id":"","label":"Block size","localized":"","hint":"","ui":"script_nudenet"},
{"id":"","label":"Banned words","localized":"","hint":"","ui":"script_nudenet"},
- {"id":"","label":"Blur","localized":"","hint":"","ui":"img2img"},
+ {"id":"","label":"Blur","localized":"","hint":"Softens the mask edge with a Gaussian blur so the boundary between masked and unmasked regions blends gradually instead of cutting hard. Reduces visible seams at the mask edge after generation. Combine with a small Dilate to push the soft transition just outside the original mask. Sigma scales with image size: at value 0.05 on a 1024px image the blur radius is roughly 13 pixels. Set to 0 to disable. Default 0.","ui":"img2img"},
{"id":"","label":"Batch input directory","localized":"","hint":"","ui":"img2img"},
{"id":"","label":"Batch output directory","localized":"","hint":"","ui":"img2img"},
{"id":"","label":"Batch mask directory","localized":"","hint":"","ui":"img2img"},
@@ -210,7 +210,7 @@
{"id":"","label":"Copy","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Composite","localized":"","hint":"","ui":"img2img"},
{"id":"control_params_elements","label":"Control","localized":"","hint":"Create image with full guidance","ui":"control"},
- {"id":"","label":"ControlNet","localized":"","hint":"ControlNet is an advanced guidance model","ui":"control"},
+ {"id":"","label":"ControlNet","localized":"","hint":"ControlNet is an advanced guidance model","ui":"control"},
{"id":"caption_tab_controls","label":"Controls","localized":"","hint":"","ui":"caption"},
{"id":"","label":"CaptionCaption","localized":"","hint":"","ui":"caption"},
{"id":"btn_console","label":"Console","localized":"","hint":""},
@@ -234,14 +234,14 @@
{"id":"","label":"Create Video","localized":"","hint":"","ui":"extras"},
{"id":"","label":"ChronoEdit","localized":"","hint":"","ui":"settings_model_options"},
{"id":"","label":"Cross Attention","localized":"","hint":"","ui":"settings_cuda"},
- {"id":"","label":"CLiP Skip","localized":"","hint":"Early stopping parameter for CLIP model; 1 is stop at last layer as usual, 2 is stop at penultimate layer, etc","ui":"settings_advanced"},
+ {"id":"","label":"CLiP Skip","localized":"","hint":"Early stopping parameter for the CLiP text encoder; 1 is stop at last layer as usual, 2 is stop at penultimate layer, etc","ui":"settings_advanced"},
{"id":"","label":"Cache-DiT","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"CFG-Zero","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"Cache folders","localized":"","hint":"","ui":"settings_system-paths"},
{"id":"","label":"Custom model loader","localized":"","hint":"","ui":"models_loader_tab"},
{"id":"","label":"Client log","localized":"","hint":""},
- {"id":"","label":"CLIP Analysis","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"Context","localized":"","hint":"","ui":"txt2img"},
+ {"id":"","label":"CLiP Analysis","localized":"","hint":"Detailed analysis output from OpenCLiP, listing the matched medium, artist, movement, trending, and flavor terms. Populated when you click the Analyze button next to the OpenCLiP Caption button.","ui":"caption"},
+ {"id":"","label":"Context","localized":"","hint":"Behavior of the Context aware resize Mode (no effect with any other Mode).Add : extend the image by inserting new pixels along smooth, featureless paths (like sky or plain backgrounds), avoiding detailed regions.Remove : shrink the image by removing pixels along the same low-detail paths.Forward : examine what the image will look like after each seam is added or removed, picking the paths that minimize visible damage. Slower but higher quality.Backward : pick paths based on existing pixel contrast in the image. Faster, classic seam-carving.","ui":"resize"},
{"id":"","label":"Contrast","localized":"","hint":"Adjusts the difference between light and dark areas. Positive values increase contrast, making darks darker and lights brighter. Negative values flatten the tonal range toward a more uniform appearance.","ui":"txt2img"},
{"id":"","label":"Color temp","localized":"","hint":"Shifts color temperature in Kelvin. Lower values (e.g., 2000K) produce a warm, amber tone. Higher values (e.g., 12000K) produce a cool, bluish tone. Default 6500K is neutral daylight. Works by scaling R/G/B channels to simulate the target white point.","ui":"txt2img"},
{"id":"","label":"CLAHE clip","localized":"","hint":"Clip limit for Contrast Limited Adaptive Histogram Equalization. Higher values allow more local contrast enhancement, which brings out detail in flat regions. Set to 0 to disable. Typical values are 1.0–3.0. Very high values can introduce noise amplification.","ui":"txt2img"},
@@ -249,7 +249,7 @@
{"id":"","label":"Correction mode","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Crop to portrait","localized":"","hint":"Crop input image to portrait-only before using it as IP adapter input","ui":"txt2img"},
{"id":"","label":"Concept Tokens","localized":"","hint":"","ui":"script_consistory"},
- {"id":"","label":"Colormap","localized":"","hint":"","ui":"script_daam"},
+ {"id":"","label":"Colormap","localized":"","hint":"OpenCV color palette used to visualize the mask or heatmap overlay. For control masks, this is the palette applied when Preview is set to Color or Composite. Pick one that contrasts well with the input image so the overlay stays readable. Default pink (control mask), jet (DAAM script).","ui":"script_daam"},
{"id":"","label":"Cosine scale 1","localized":"","hint":"","ui":"script_demofusion"},
{"id":"","label":"Cosine scale 2","localized":"","hint":"","ui":"script_demofusion"},
{"id":"","label":"Cosine scale 3","localized":"","hint":"","ui":"script_demofusion"},
@@ -267,9 +267,9 @@
{"id":"","label":"Control override denoise strength","localized":"","hint":"","ui":"script_flux_tools"},
{"id":"","label":"Color variation","localized":"","hint":"","ui":"script_outpainting"},
{"id":"","label":"Change rate","localized":"","hint":"","ui":"script_video"},
- {"id":"","label":"Context after","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Context mask","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Control only","localized":"","hint":"This uses only the Control input below as the source for any ControlNet or IP Adapter type tasks based on any of our various options.","ui":"control"},
+ {"id":"","label":"Context after","localized":"","hint":"Behavior of the Context aware resize Mode applied to the output image after the model finishes generating (Post sub-tab in the Size accordion; no effect with any other Mode).Add : extend the image by inserting new pixels along smooth, featureless paths (like sky or plain backgrounds), avoiding detailed regions.Remove : shrink the image by removing pixels along the same low-detail paths.Forward : examine what the image will look like after each seam is added or removed, picking the paths that minimize visible damage. Slower but higher quality.Backward : pick paths based on existing pixel contrast in the image. Faster, classic seam-carving.","ui":"control"},
+ {"id":"","label":"Context mask","localized":"","hint":"Behavior of the Context aware resize Mode applied to the input mask image (used for inpainting, outpainting, or control masks; Mask sub-tab in the Size accordion; no effect with any other Mode).Add : extend the image by inserting new pixels along smooth, featureless paths (like sky or plain backgrounds), avoiding detailed regions.Remove : shrink the image by removing pixels along the same low-detail paths.Forward : examine what the image will look like after each seam is added or removed, picking the paths that minimize visible damage. Slower but higher quality.Backward : pick paths based on existing pixel contrast in the image. Faster, classic seam-carving.","ui":"control"},
+ {"id":"","label":"Control only","localized":"","hint":"This uses only the Control input below as the source for any ControlNet or IP Adapter type tasks based on any of our various options.","ui":"control"},
{"id":"","label":"CN Mode","localized":"","hint":"","ui":"control"},
{"id":"","label":"CN Strength","localized":"","hint":"","ui":"control"},
{"id":"","label":"CN Start","localized":"","hint":"","ui":"control"},
@@ -280,14 +280,14 @@
{"id":"","label":"Coarse","localized":"","hint":"","ui":"control"},
{"id":"","label":"Color map","localized":"","hint":"","ui":"control"},
{"id":"","label":"Crop to fit","localized":"","hint":"If the dimensions of your source image (e.g. 512x510) deviate from your target dimensions (e.g. 1024x768) this function will fit your upscaled image into your target size image. Excess will be cropped","ui":"extras"},
- {"id":"","label":"CLiP Model","localized":"","hint":"CLIP model used for image-text similarity matching. Larger models (ViT-L, ViT-H) are more accurate but slower and use more VRAM.","ui":"caption"},
+ {"id":"","label":"CLiP Model","localized":"","hint":"CLiP model used for image-text similarity matching. Larger models (ViT-L, ViT-H) are more accurate but slower and use more VRAM.","ui":"caption"},
{"id":"","label":"Caption Model","localized":"","hint":"BLIP model used to generate the initial image caption. The caption model describes the image content which CLiP then enriches with style and flavor terms.","ui":"caption"},
- {"id":"","label":"clip: max length","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"clip: chunk size","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"clip: min flavors","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"clip: max flavors","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"clip: intermediates","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"clip: num beams","localized":"","hint":"","ui":"caption"},
+ {"id":"","label":"clip: max length","localized":"Max Length","hint":"Maximum number of tokens in the generated caption. Higher values allow longer, more descriptive captions; lower values produce shorter ones.","ui":"caption"},
+ {"id":"","label":"clip: chunk size","localized":"Chunk Size","hint":"Batch size for processing description candidates (flavors). Higher values speed up interrogation but increase VRAM usage.","ui":"caption"},
+ {"id":"","label":"clip: min flavors","localized":"Min Flavors","hint":"Minimum number of descriptive tags (flavors) to keep in the final prompt.","ui":"caption"},
+ {"id":"","label":"clip: max flavors","localized":"Max Flavors","hint":"Maximum number of descriptive tags (flavors) to keep in the final prompt.","ui":"caption"},
+ {"id":"","label":"clip: intermediates","localized":"Intermediates","hint":"Size of the intermediate candidate pool when matching image features to descriptive tags (flavors). From this pool, the final tags are selected based on Min/Max Flavors. Higher values may improve quality but are slower.","ui":"caption"},
+ {"id":"","label":"clip: num beams","localized":"CLiP Num Beams","hint":"Number of beams for beam search during caption generation. Higher values search more possibilities but are slower. Set to 1 to disable beam search.","ui":"caption"},
{"id":"","label":"Character threshold","localized":"","hint":"Confidence threshold for character-specific tags (e.g., character names, specific traits). Only tags with confidence above this threshold are included. Higher values are more selective, lower values include more potential matches. Not supported by DeepBooru models.","ui":"caption"},
{"id":"","label":"Cross-attention","localized":"","hint":"","ui":"component-8779"},
{"id":"","label":"cpu","localized":"","hint":"Uses cpu and RAM only: slowest but least likely to OOM","ui":"settings_sd"},
@@ -343,8 +343,8 @@
{"id":"","label":"Control settings","localized":"","hint":"","ui":"control"},
{"id":"","label":"Canny","localized":"","hint":"","ui":"control"},
{"id":"","label":"Condition","localized":"","hint":"","ui":"video"},
- {"id":"","label":"Caption: Advanced Options","localized":"","hint":"","ui":"caption"},
- {"id":"","label":"Caption: Batch","localized":"","hint":"","ui":"caption"},
+ {"id":"","label":"Caption: Advanced Options","localized":"","hint":"Advanced configuration options for caption generation. Sampling parameters, length limits, and decoding behavior for the active backend (VLM, CLiP, or Tagger).","ui":"caption"},
+ {"id":"","label":"Caption: Batch","localized":"","hint":"Process multiple images in a batch using the active caption backend. Captions are saved alongside the source images as .txt sidecar files when Save Caption Files is enabled.","ui":"caption"},
{"id":"","label":"Control elements","localized":"","hint":"Control elements are advanced models that can guide generation towards desired outcome","ui":"tab_control"}
],
"d": [
@@ -365,22 +365,22 @@
{"id":"","label":"Download model from huggingface","localized":"","hint":"","ui":"models_huggingface_tab"},
{"id":"","label":"Dropdown","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"dynamic","localized":"","hint":"Dynamic shifting automatically adjusts the denoising schedule based on your image resolution. The scheduler interpolates between base_shift and max_shift based on actual image resolution. Enabling disables static Flow shift.","ui":"txt2img"},
- {"id":"","label":"Detailer models","localized":"","hint":"Select detection models to use for detailing","ui":"txt2img"},
+ {"id":"","label":"Detailer models","localized":"","hint":"YOLO detection models used to find regions to re-render. Multiple models can be selected and they run in sequence. Models live in models/yolo. Filename hints at target: face-* detects faces, eyes-* detects eyes, hand-* detects hands, person-* detects whole subjects, and so on. Models with -seg in the name produce a precise segmentation outline (used when Use segmentation is on); the rest produce only bounding boxes. Per-model overrides can be appended with colon syntax, for example face-yolo8n:conf=0.5:strength=0.4.","ui":"txt2img"},
{"id":"","label":"Detailer list","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"Detailer classes","localized":"","hint":"Specify specific classes to use if selected detailer model is a multi-class model","ui":"txt2img"},
- {"id":"","label":"Detailer prompt","localized":"","hint":"Use separate prompt for detailer. If not present, it will use primary prompt","ui":"txt2img"},
- {"id":"","label":"Detailer negative prompt","localized":"","hint":"Use separate negative prompt for detailer. If not present, it will use primary negative prompt","ui":"txt2img"},
- {"id":"","label":"Detailer steps","localized":"","hint":"Number of steps to run for detailer process","ui":"txt2img"},
- {"id":"","label":"Detailer strength","localized":"","hint":"Denoising strength of detailer process","ui":"txt2img"},
- {"id":"","label":"Detailer resolution","localized":"","hint":"","ui":"txt2img"},
+ {"id":"","label":"Detailer classes","localized":"","hint":"Comma-separated list of class names to keep when the selected detailer model is multi-class (e.g., a YOLO model that detects faces, eyes, and hands all in one file). Only detections matching these labels are processed; everything else is dropped. Leave empty to accept all classes. Names must match the model's class names exactly (case-insensitive). Single-class models like a face-only detector ignore this field.","ui":"txt2img"},
+ {"id":"","label":"Detailer prompt","localized":"","hint":"Optional dedicated prompt for the detailer pass. Leave empty to inherit the main prompt. Useful for steering the inpaint differently from the rest of the image: a face detailer can use just portrait, sharp eyes, detailed skin while the main prompt covers the full scene. The placeholder [PROMPT] (or [prompt]) is replaced with the original main prompt, so you can append to it: [PROMPT], detailed face.","ui":"txt2img"},
+ {"id":"","label":"Detailer negative prompt","localized":"","hint":"Optional dedicated negative prompt for the detailer pass. Leave empty to inherit the main negative prompt. Same [PROMPT] / [prompt] placeholder behavior as the positive detailer prompt: it expands to the original main negative prompt.","ui":"txt2img"},
+ {"id":"","label":"Detailer steps","localized":"","hint":"Number of sampling steps used for each detailer inpaint pass. Independent of the main generation steps. Higher values give cleaner detail but cost more time per detected region. Set to 0 to inherit the main generation step count. Default 10.","ui":"txt2img"},
+ {"id":"","label":"Detailer strength","localized":"","hint":"Denoising strength of the detailer inpaint pass. Higher values regenerate more aggressively (more change to the detected region, more reliance on the prompt). Lower values stay closer to the original detection, only refining detail. Typical range 0.2 to 0.5: enough to fix distortions without losing identity. Above 0.7 the face/object can drift noticeably from the original. Set to 0 to skip the detailer pass entirely. Default 0.30.","ui":"txt2img"},
+ {"id":"","label":"Detailer resolution","localized":"","hint":"Working resolution for the detailer inpaint pass. Each detected region is cropped (with Edge padding ) and resized to this resolution before inpainting. Higher values give finer detail in the regenerated region but use more VRAM and time per detection. Match the model's native resolution for best results: 1024 for SDXL /SD3 /Flux , 512 for SD 1.5 . Default 1024.","ui":"txt2img"},
{"id":"","label":"Denoising batch size","localized":"","hint":"","ui":"script_demofusion"},
{"id":"","label":"Dilate tau","localized":"","hint":"","ui":"script_freescale"},
{"id":"","label":"Draw legend","localized":"","hint":"","ui":"script_xyz_grid_script"},
- {"id":"","label":"Denoising strength","localized":"","hint":"Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image. With values below 1.0, processing will take less steps than the Sampling Steps slider specifies","ui":"img2img"},
+ {"id":"","label":"Denoising strength","localized":"","hint":"Strength of img2img modification when an init image is supplied. Higher values move further from the init image and rely more on the prompt; lower values stay closer to the original. At 0.0 the init image passes through unchanged. At 1.0 the model builds a fresh image from scratch and effectively ignores the init image. Effect on step count is model-dependent: - SD 1.5 and SDXL : the configured Steps value is honored as the actual loop count; strength only controls how much noise is added to the init latent. - Flux, SD3, Hunyuan, Sana, Qwen and other DiT models : loop count is reduced proportionally; with strength 0.5 and 30 steps, only ~15 actually run. In the Images tab this only takes effect when Use init image is set to one of the init modes; with No: Control only it is ignored. Default 0.30.","ui":"img2img"},
{"id":"","label":"Denoise start","localized":"","hint":"Override denoise strength by stating how early base model should finish and when refiner should start. Only applicable to refiner usage. If set to 0 or 1, denoising strength will be used","ui":"img2img"},
{"id":"","label":"down","localized":"","hint":"","ui":"script_outpainting"},
{"id":"","label":"Decode chunks","localized":"","hint":"","ui":"script_video"},
- {"id":"","label":"Dilate","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Dilate","localized":"","hint":"Expands the masked area outward by growing each masked pixel into its neighborhood. Useful for catching the edges around an object that the mask missed, or for giving the model more breathing room around the region being modified so the new content can blend with surrounding context. Kernel size scales with image size: at value 0.05 on a 1024px image the dilation reaches roughly 13 pixels in each direction. Set to 0 to disable. Default 0.","ui":"control"},
{"id":"","label":"Depth and normal","localized":"","hint":"","ui":"control"},
{"id":"","label":"Distance threshold","localized":"","hint":"","ui":"control"},
{"id":"","label":"Depth threshold","localized":"","hint":"","ui":"control"},
@@ -440,9 +440,9 @@
{"id":"","label":"Effects","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Enable LayerSkipConfig","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Enable refine pass","localized":"","hint":"Use a similar process as image to image to upscale and/or add detail to the final image. Optionally uses refiner model to enhance image details.","ui":"txt2img"},
- {"id":"","label":"Enable detailer pass","localized":"","hint":"Detect target objects such as face and reprocess it at higher resolution","ui":"txt2img"},
- {"id":"","label":"Edge padding","localized":"","hint":"Expand edge of masked area by this percentage","ui":"txt2img"},
- {"id":"","label":"Edge blur","localized":"","hint":"Blur edge of masked area by this percentage","ui":"txt2img"},
+ {"id":"","label":"Enable detailer pass","localized":"","hint":"Runs an automatic touch-up pass after generation: a YOLO detector finds target regions (faces, eyes, hands, persons, etc.) and each detected region is re-rendered with inpaint at the configured detailer resolution. Useful for fixing distorted faces or hands at low base resolutions, sharpening eye detail, or adding a second-pass refinement to specific subjects. Default off.","ui":"txt2img"},
+ {"id":"","label":"Edge padding","localized":"","hint":"Pixels added around each detection's bounding box when cropping the region for inpaint. Padding gives the inpaint pass surrounding context so the regenerated content can blend smoothly with the rest of the image. Too little causes hard seams; too much wastes resolution on areas that won't change. Default 20.","ui":"txt2img"},
+ {"id":"","label":"Edge blur","localized":"","hint":"Pixel radius of the Gaussian blur applied to the inpaint mask edge. Softens the boundary between the regenerated region and the rest of the image so the paste-back blends instead of cutting hard. Set to 0 to disable. Default 10.","ui":"txt2img"},
{"id":"","label":"End","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"ETA","localized":"","hint":"","ui":"script_apg"},
{"id":"","label":"Enable FreeU","localized":"","hint":"","ui":"script_consistory"},
@@ -452,7 +452,7 @@
{"id":"","label":"Enhanced prompt","localized":"","hint":"The enhanced prompt output from the LLM","ui":"script_prompt_enhance"},
{"id":"","label":"Edit start","localized":"","hint":"","ui":"script_ledits"},
{"id":"","label":"Edit stop","localized":"","hint":"","ui":"script_ledits"},
- {"id":"","label":"Erode","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Erode","localized":"","hint":"Shrinks the masked area inward by removing pixels along the edge. Useful for cleaning up speckle noise from auto-segmentation, or for pulling the mask back from object boundaries to avoid the model bleeding outside the intended region. Kernel size scales with image size: at value 0.05 on a 1024px image the erosion reaches roughly 13 pixels in each direction. Set to 0 to disable. Default 0.","ui":"control"},
{"id":"","label":"edge","localized":"","hint":"","ui":"control"},
{"id":"","label":"Ensemble size","localized":"","hint":"","ui":"control"},
{"id":"","label":"Enable","localized":"","hint":"","ui":"video"},
@@ -490,7 +490,7 @@
{"id":"","label":"Faster Cache","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"Folders","localized":"","hint":"","ui":"settings_saving-paths"},
{"id":"","label":"Fetch model preview metadata","localized":"","hint":"","ui":"models_metadata_tab"},
- {"id":"","label":"Flow shift","localized":"","hint":"Shift value for flowmatching models. Controls the distribution of denoising steps. Values: - >1.0: allocate more steps to early denoising (better structure) -<1.0: allocate more steps to late denoising (better fine details) - 1.0: balanced schedule Most flowmatching models use the value of 3 as default. Effectively inactive if dynamic shift is enabled.","ui":"txt2img"},
+ {"id":"","label":"Flow shift","localized":"","hint":"Shift value for flowmatching models. Controls the distribution of denoising steps. Values: - >1.0 : allocate more steps to early denoising (better structure) - <1.0 : allocate more steps to late denoising (better fine details) - 1.0 : balanced schedule Most flowmatching models use the value of 3 as default. Effectively inactive if dynamic shift is enabled.","ui":"txt2img"},
{"id":"","label":"FDG scales","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"FDG weights","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"FDG rescale space","localized":"","hint":"","ui":"txt2img"},
@@ -622,9 +622,9 @@
{"id":"","label":"Grid Options","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"Grids","localized":"","hint":"","ui":"settings_saving-paths"},
{"id":"","label":"Guider","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"Guidance scale","localized":"","hint":"Classifier Free Guidance scale: how strongly the image should conform to prompt. Lower values produce more creative results, higher values make it follow the prompt more strictly; recommended values between 5-10","ui":"txt2img"},
- {"id":"","label":"Guidance end","localized":"","hint":"Ends the effect of CFG and PAG early: A value of 1 acts as normal, 0.5 stops guidance at 50% of steps","ui":"txt2img"},
- {"id":"","label":"Guidance rescale","localized":"","hint":"Rescale guidance to avoid overexposed images at higher guidance values","ui":"txt2img"},
+ {"id":"","label":"Guidance scale","localized":"","hint":"Classifier-Free Guidance scale. How strongly the image should conform to the prompt. Lower values produce more creative, loosely-prompted results; higher values follow the prompt more strictly but can oversaturate or burn out at very high values. Recommended values vary by architecture: 5-10 for SDXL /SD1.x , 3-5 for Flux and SD3 , 7-10 for video models. Check the model card if unsure. Set to 1 (the slider's minimum) to disable guidance entirely. The model then runs only the conditional prediction with no negative-prompt steering.","ui":"txt2img"},
+ {"id":"","label":"Guidance end","localized":"","hint":"Ends guidance early. The remaining denoising steps run unguided, which can speed up inference and produce slightly softer, less prompt-locked results. Applied independently to each pipeline pass (base, HiRes, refiner) against that pass's own step count. Example: 0.5 stops guidance at 50% of steps; 0.8 stops at 80%. Affects Guidance scale and Refine guidance on all models, and Attention guidance on the PAG path only (SD 1.5 and SDXL ). Has no effect on the true_cfg_scale path that Attention guidance uses for Flux , QwenImage , HiDream , Hunyuan Video , and other flow-matching models. Set to 1 to keep guidance active for the entire denoising process. 1 (no early end) by default.","ui":"txt2img"},
+ {"id":"","label":"Guidance rescale","localized":"","hint":"Rescales the guided noise prediction to avoid the oversaturated, washed-out colors that high Guidance scale values can produce. Useful when running with Guidance scale above 10 or when colors look blown out. Mild values (0.5-0.7) usually fix the issue without affecting prompt adherence. Set to 0 to disable rescaling. Disabled by default.","ui":"txt2img"},
{"id":"","label":"Gamma","localized":"","hint":"Non-linear brightness curve adjustment. Values below 1.0 brighten midtones and shadows while preserving highlights. Values above 1.0 darken midtones and shadows. Default is 1.0 (no change). Unlike brightness, gamma reshapes the tonal curve rather than shifting it uniformly.","ui":"txt2img"},
{"id":"","label":"Grain","localized":"","hint":"Adds film-like noise to the image. Higher values produce more visible grain, simulating analog film texture. Applied as random noise blended into the final image. Set to 0 to disable.","ui":"txt2img"},
{"id":"","label":"Grid margins","localized":"","hint":"","ui":"script_prompt_matrix"},
@@ -633,7 +633,7 @@
{"id":"","label":"Guidance start","localized":"","hint":"","ui":"script_slg"},
{"id":"","label":"Guidance stop","localized":"","hint":"","ui":"script_slg"},
{"id":"","label":"Gate step","localized":"","hint":"","ui":"script_t-gate"},
- {"id":"","label":"Guess mode","localized":"","hint":"Removes the requirement to supply a prompt to a ControlNet. It forces Controlnet encoder to do it's 'best guess' based on the contents of the input control map.","ui":"control"},
+ {"id":"","label":"Guess mode","localized":"","hint":"Removes the requirement to supply a prompt to a ControlNet . It forces ControlNet encoder to do its 'best guess' based on the contents of the input control map.","ui":"control"},
{"id":"","label":"gradient","localized":"","hint":"","ui":"control"},
{"id":"","label":"Gamma corrected","localized":"","hint":"","ui":"control"},
{"id":"","label":"General threshold","localized":"","hint":"Confidence threshold for general tags (e.g., objects, actions, settings). Only tags with confidence above this threshold are included in the output. Higher values are more selective (fewer tags), lower values include more tags.","ui":"caption"},
@@ -661,15 +661,15 @@
{"id":"","label":"HiDream","localized":"","hint":"","ui":"settings_model_options"},
{"id":"","label":"HyperTile","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"HiDiffusion","localized":"","hint":"HiDiffusion allows creation of high-resolution images using your standard models without duplicates/distortions and improved performance","ui":"settings_advanced"},
- {"id":"","label":"Height","localized":"","hint":"Image height","ui":"txt2img"},
+ {"id":"","label":"Height","localized":"","hint":"Target height of the output image in pixels. For generation, this sets the resolution the model produces. For resize and upscale operations, this is the height the input is fitted to. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"txt2img"},
{"id":"","label":"HiRes steps","localized":"","hint":"Number of sampling steps for upscaled picture. If 0, uses same as for original","ui":"txt2img"},
{"id":"","label":"Hue","localized":"","hint":"Rotates all colors around the color wheel. Small values produce subtle color shifts, while higher values cycle through the full spectrum. Useful for creative color effects or correcting unwanted color casts.","ui":"txt2img"},
{"id":"","label":"Highlights","localized":"","hint":"Adjusts the brightness of highlight (bright) regions. Positive values brighten highlights, negative values pull them down. Operates on the L channel in Lab color space using a luminance-weighted mask, leaving shadows and midtones largely unaffected.","ui":"txt2img"},
{"id":"","label":"Highlights tint","localized":"","hint":"Color to blend into highlight regions for split toning. Works together with Shadows tint and Split tone balance to create cinematic color grading looks. Default white (#ffffff) applies no tint.","ui":"txt2img"},
{"id":"","label":"HDR range","localized":"","hint":"","ui":"script_hdr"},
{"id":"","label":"HQ init latents","localized":"","hint":"","ui":"script_instantir"},
- {"id":"","label":"Height after","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Height mask","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Height after","localized":"","hint":"Target height of the output image in pixels, applied after the model finishes generating (Post sub-tab in the Size accordion). Use this to upscale or downscale the final image before saving. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"control"},
+ {"id":"","label":"Height mask","localized":"","hint":"Target height of the input mask image in pixels (Mask sub-tab in the Size accordion). The mask is used for inpainting, outpainting, or as a control mask, and is resized so it aligns with the processing resolution. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"control"},
{"id":"","label":"Hires use control","localized":"","hint":"","ui":"control"},
{"id":"","label":"Hands","localized":"","hint":"","ui":"control"},
{"id":"","label":"High threshold","localized":"","hint":"","ui":"control"},
@@ -700,7 +700,6 @@
"i": [
{"id":"control_nav","label":"Images","localized":"","hint":"Create images Unified interface Supports T2I and I2I With optional control guidance"},
{"id":"img2img_nav","label":"I2I","localized":"","hint":"Create image from image Legacy interface that mimics original image-to-image interface and behavior"},
- {"id":"img2img_results_input_mobile","label":"Input","localized":"","hint":"Show/hide selection of input media used to guide generation","ui":"img2img"},
{"id":"","label":"Image","localized":"","hint":"Create image from image","ui":"img2img"},
{"id":"","label":"Inpaint","localized":"","hint":"","ui":"img2img"},
{"id":"control_params_mask","label":"Inputs","localized":"","hint":"Settings related to Input images","ui":"control"},
@@ -711,13 +710,13 @@
{"id":"","label":"Image Paths","localized":"","hint":"Settings related to image filenames, and output directories"},
{"id":"","label":"Image Metadata","localized":"","hint":"Settings related to handling of metadata that is created with generated images"},
{"id":"","label":"IP Adapters","localized":"","hint":"IP adapters are plugin models that can guide generation towards desired outcome","ui":"txt2img"},
- {"id":"","label":"Input Media","localized":"","hint":"Add input image to be used for image-to-image, inpaint or control processing","ui":"control"},
+ {"id":"","label":"Input","localized":"","hint":"Add input image to be used for image-to-image, inpaint or control processing Click to minimize/maximize","ui":"control"},
{"id":"","label":"Input Image","localized":"","hint":"","ui":"caption"},
{"id":"","label":"IPEX","localized":"","hint":"","ui":"settings_backends"},
{"id":"","label":"Image Gallery","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"Intermediate Image Saving","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"Initial seed","localized":"","hint":"A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result","ui":"txt2img"},
- {"id":"","label":"Include detections","localized":"","hint":"Include original image with detected areas marked","ui":"txt2img"},
+ {"id":"","label":"Include detections","localized":"","hint":"Adds an annotated debug image to the output gallery showing each detected region's bounding box, label, and confidence score, plus a translucent mask overlay. Useful for tuning Min confidence , Min size /Max size , and class filters: you can see exactly what was detected before the inpaint pass touched the image. Default off.","ui":"txt2img"},
{"id":"","label":"IY model","localized":"","hint":"","ui":"script_infiniteyou"},
{"id":"","label":"IY scale","localized":"","hint":"","ui":"script_infiniteyou"},
{"id":"","label":"IY start","localized":"","hint":"","ui":"script_infiniteyou"},
@@ -731,8 +730,8 @@
{"id":"","label":"Include images","localized":"","hint":"","ui":"script_xyz_grid_script"},
{"id":"","label":"invert","localized":"","hint":"","ui":"img2img"},
{"id":"","label":"Init image same as control","localized":"","hint":"Will additionally treat any image placed into the Control input window as a source for img2img type tasks, an image to modify for example.","ui":"control"},
- {"id":"","label":"Inpaint masked only","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Invert mask","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Inpaint masked only","localized":"","hint":"Crop the masked region, denoise it at full resolution, then paste the result back into the original image. Best for small detail edits where you want maximum quality on the masked area without spending compute denoising the rest of the image. Detail in unmasked regions stays untouched. Tradeoff: the model only sees the cropped region, so it loses global context. The inpainted content may not match the surrounding scene's lighting, perspective, or style, and visible seams can appear at the crop boundary. Mitigate with Dilate + Blur on the mask, or disable this option to denoise the full image together. When off, the whole image is denoised at the generation resolution and the unmasked area is restored from the original via the mask blend, which preserves global coherence at the cost of detail in the masked region. Default off.","ui":"control"},
+ {"id":"","label":"Invert mask","localized":"","hint":"Swaps which area is treated as masked. Useful when you have painted the region to preserve instead of the region to modify : enable this to flip the interpretation without redoing the mask. Default off.","ui":"control"},
{"id":"","label":"IOU","localized":"","hint":"","ui":"control"},
{"id":"","label":"Init strength","localized":"","hint":"","ui":"video"},
{"id":"","label":"Input directory","localized":"","hint":"Folder where the images are that you want to process","ui":"extras"},
@@ -754,7 +753,7 @@
{"id":"","label":"Image resize algorithm","localized":"","hint":"","ui":"settings_postprocessing"},
{"id":"","label":"Image repeats per epoch","localized":"","hint":"","ui":"settings_legacy_options"},
{"id":"","label":"Interpolation Method","localized":"","hint":"","ui":"models_merge_tab"},
- {"id":"","label":"In Blocks","localized":"","hint":"Downsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)","ui":"component-5674"},
+ {"id":"","label":"In Blocks","localized":"","hint":"Downsampling Blocks of the UNet (12 values for SD1.5 , 9 values for SDXL )","ui":"component-5674"},
{"id":"","label":"Input model","localized":"","hint":"","ui":"models_replace_tab"},
{"id":"","label":"Info object","localized":"","hint":"","ui":"component-8779"}
],
@@ -765,12 +764,13 @@
{"id":"","label":"Keep Thinking Trace","localized":"","hint":"Include the model's reasoning process in the final output. Useful for understanding how the model arrived at its answer. Only works with models that support thinking mode.","ui":"script_prompt_enhance"},
{"id":"","label":"Keep Prefill","localized":"","hint":"Include the prefill text at the beginning of the final output. If disabled, the prefill text used to guide the model is removed from the result.","ui":"script_prompt_enhance"},
{"id":"","label":"Keep aspect ratio","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Keep @ on artist insert","localized":"","hint":"Type @ in the prompt to filter autocomplete to artist tags only. This setting controls only what gets inserted on accept; the @ filter works for every model.Enable for models that require the @ prefix in the prompt itself, e.g. Anima . Inserts as @artist name with underscores converted to spaces.Disable for booru-trained models that take plain artist tags, e.g. SDXL , Pony , Illustrious , NoobAI . The typed @ is consumed and the artist name is inserted as a normal tag.","ui":"script_autocomplete"}
+ {"id":"","label":"Keep @ on artist insert","localized":"","hint":"Type @ in the prompt to filter autocomplete to artist tags only. This setting controls only what gets inserted on accept; the @ filter works for every model. Underscore handling is controlled by Keep underscores .Enable for models that require the @ prefix in the prompt itself, e.g. Anima . Inserts as @artist name.Disable for booru-trained models that take plain artist tags, e.g. SDXL , Pony , Illustrious , NoobAI . The typed @ is consumed and the artist name is inserted as a normal tag.","ui":"script_autocomplete"},
+ {"id":"","label":"Keep underscores","localized":"","hint":"Keep underscore characters when inserting tags from autocomplete. Applies to both ordinary tags and artist insertions (the @ trigger). Embedding names always preserve their underscores regardless of this setting.Enable when your model is sensitive to the underscored form of booru tags. The tag long_hair displays and inserts as long_hair.Disable (default) to convert underscores to spaces, matching the prompting style of most modern checkpoints. The tag long_hair displays and inserts as long hair.","ui":"script_autocomplete"}
],
"l": [
{"id":"prompt_enhance_load","label":"Load model","localized":"","hint":"","ui":"script_prompt_enhance"},
{"id":"prompt_enhance_custom_load","label":"Load custom model","localized":"","hint":"Load a custom model with the specified configuration","ui":"script_prompt_enhance"},
- {"id":"control_mask_remove","label":"LaMa Remove","localized":"","hint":"","ui":"control"},
+ {"id":"control_mask_remove","label":"LaMa Remove","localized":"","hint":"Removes the masked region using LaMa, a lightweight inpainting model that fills the area with content extrapolated from the surroundings. Useful for cleanup tasks like erasing watermarks, removing unwanted objects, or generating a clean plate before running a full diffusion pass. Runs the configured mask pipeline (auto-segment, dilate, erode, blur, invert) first, then passes the resulting mask to LaMa. Result is written to the output panel. Model is downloaded on first use.","ui":"control"},
{"id":"","label":"Lite","localized":"","hint":"","ui":"control"},
{"id":"video_params_ltx","label":"LTXVideo","localized":"","hint":"","ui":"video"},
{"id":"vlm_load","label":"Load","localized":"","hint":"","ui":"caption"},
@@ -790,7 +790,7 @@
{"id":"","label":"List all locally available models","localized":"","hint":"","ui":"models_list_tab"},
{"id":"","label":"Last Generate","localized":"","hint":""},
{"id":"","label":"LUT","localized":"","hint":"Look-Up Table color grading section. Upload a .cube LUT file to apply professional color grading presets. LUTs remap colors according to a predefined 3D color transform, commonly used in film and photography for consistent color looks.","ui":"txt2img"},
- {"id":"","label":"low order","localized":"","hint":"","ui":"txt2img"},
+ {"id":"","label":"low order","localized":"","hint":"Forces multistep solvers to fall back to a lower-order step during the last few denoising iterations. Higher-order solvers can become numerically unstable as sigma approaches zero, so the lower-order tail produces a cleaner, more stable final image. Applies only to multistep families (DPM++ , UniPC , DEIS , SA Solver , DC Solver , ER-SDE ). Single-step samplers such as DDIM , Euler , and Euler a ignore this option. Recommended to leave on. Disabling can occasionally give slightly sharper output but risks artifacts on the final steps. Enabled by default.","ui":"txt2img"},
{"id":"","label":"LSC layer indices","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"LSC fully qualified name","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"LSC skip attention blocks","localized":"","hint":"","ui":"txt2img"},
@@ -831,7 +831,7 @@
{"id":"","label":"LTX enable refine","localized":"","hint":"","ui":"video"},
{"id":"","label":"LTX refine strength","localized":"","hint":"","ui":"video"},
{"id":"","label":"LTX decode timestep","localized":"","hint":"","ui":"video"},
- {"id":"","label":"LTX enable audio","localized":"","hint":"","ui":"video"},
+ {"id":"","label":"LTX save audio","localized":"","hint":"LTX-2 audio-capable models always generate audio from the same prompt as video; this toggle controls whether the audio track is included in the saved video file","ui":"video"},
{"id":"","label":"Loop","localized":"","hint":"","ui":"extras"},
{"id":"","label":"Local directory name","localized":"","hint":"Directory where to install extension, leave blank for default","ui":"component-8746"},
{"id":"","label":"Libs","localized":"","hint":"","ui":"component-8779"},
@@ -889,20 +889,21 @@
{"id":"","label":"Mobile","localized":"","hint":"","ui":"settings_ui"},
{"id":"","label":"Merge multiple models","localized":"","hint":"","ui":"models_merge_tab"},
{"id":"","label":"Max shift","localized":"","hint":"Maximum shift value for high resolutions when using dynamic shifting.","ui":"txt2img"},
- {"id":"","label":"Merge detailers","localized":"","hint":"Merge results from multiple detailers into single mask before running detailing process","ui":"txt2img"},
- {"id":"","label":"Max detected","localized":"","hint":"Maximum number of detected objects to run detailer on","ui":"txt2img"},
- {"id":"","label":"Min confidence","localized":"","hint":"Minimum confidence in detected item","ui":"txt2img"},
- {"id":"","label":"Max overlap","localized":"","hint":"Maximum overlap between two detected items before one is discarded","ui":"txt2img"},
- {"id":"","label":"Min size","localized":"","hint":"Minimum size of detected object as percentage of overal image","ui":"txt2img"},
- {"id":"","label":"Max size","localized":"","hint":"Maximum size of detected object as percentage of overal image","ui":"txt2img"},
+ {"id":"","label":"Merge detailers","localized":"","hint":"Combines all detections from each model into a single mask and runs one inpaint pass per model instead of one per detection. Faster when many regions are detected (e.g., a crowd scene with multiple faces): one larger inpaint pass replaces several small ones. Tradeoff: each region gets less individual attention because the model sees them all together. Best for scenes where the detected regions are similar in size and content. Default off.","ui":"txt2img"},
+ {"id":"","label":"Max detected","localized":"","hint":"Cap on how many detections per model are processed. Detections beyond this count are dropped (in detection score order, highest first). Use to keep detailer time bounded on busy scenes. Default 2.","ui":"txt2img"},
+ {"id":"","label":"Min confidence","localized":"","hint":"Minimum YOLO detection score required for a region to be processed. Higher values keep only confident detections (fewer false positives but may miss real subjects in difficult lighting). Lower values include more candidates including weak ones. Tune with Include detections on so you can see what is being kept and dropped. Default 0.6.","ui":"txt2img"},
+ {"id":"","label":"Max overlap","localized":"","hint":"IOU threshold for non-maximum suppression: if two detections overlap by more than this fraction, the lower-scoring one is dropped. Lower values are stricter (less overlap allowed; fewer duplicate detections of the same subject). Higher values let near-duplicates through, which is rarely useful. Default 0.5.","ui":"txt2img"},
+ {"id":"","label":"Min size","localized":"","hint":"Minimum detection size as a fraction of the image's shorter edge. Detections smaller than this are dropped. Use to filter out tiny background objects (e.g., faces in a crowd that aren't worth detailing). At 0.1, a face must occupy at least 10% of the image dimension to qualify. Set to 0 to disable the lower bound. Default 0.","ui":"txt2img"},
+ {"id":"","label":"Max size","localized":"","hint":"Maximum detection size as a fraction of the image's shorter edge. Detections larger than this are dropped. Use to skip cases where the detector grabs the whole image (e.g., a person detector returning a near full-frame box that the inpaint pass would just regenerate). Set to 1.0 to disable the upper bound. Default 0.75.","ui":"txt2img"},
{"id":"","label":"Midtones","localized":"","hint":"Adjusts the brightness of midtone regions. Positive values brighten midtones, negative values darken them. Targets pixels near the middle of the luminance range using a bell-shaped mask in Lab space, leaving shadows and highlights largely untouched.","ui":"txt2img"},
{"id":"","label":"Momentum","localized":"","hint":"","ui":"script_apg"},
{"id":"","label":"Mode x-axis","localized":"","hint":"","ui":"script_asymmetric_tiling"},
{"id":"","label":"Mode y-axis","localized":"","hint":"","ui":"script_asymmetric_tiling"},
{"id":"","label":"Mask Dropout","localized":"","hint":"","ui":"script_consistory"},
{"id":"","label":"Multi decoder","localized":"","hint":"","ui":"script_demofusion"},
- {"id":"","label":"Mode","localized":"","hint":"Interrogation mode.Fast : Quick caption with minimal flavor terms.Classic : Standard interrogation with balanced quality and speed.Best : Most thorough analysis, slowest but highest quality.Negative : Generate terms to use as negative prompt.","ui":"script_face"},
- {"id":"","label":"Method","localized":"","hint":"","ui":"script_video"},
+ {"id":"","label":"CLiP Mode","localized":"","hint":"OpenCLiP interrogation depth.Fast : quick caption with minimal flavor terms.Classic : standard interrogation balancing quality and speed.Best : most thorough analysis, slowest but highest quality.Negative : generate terms suitable for use as a negative prompt.","ui":"caption"},
+ {"id":"","label":"Mode","localized":"","hint":"How the input is fitted to the target resolution.None : skip resize, pass the image through unchanged.Fixed : force to target width and height, distorting aspect ratio if they differ.Crop : scale to fully cover the target then center-crop the overflow, preserving aspect ratio.Fill : scale to fit inside the target then pad the remaining space with the background color (set in Settings → Image options).Outpaint : like Fill, but the model paints new content into the padded space instead of using a solid color.Context aware : use seam-carving to add or remove pixels along smooth, featureless paths through the image (like sky or plain backgrounds), preserving the detailed regions. Behavior is controlled by the Context dropdown next to this one.","ui":"resize"},
+ {"id":"","label":"Method","localized":"","hint":"Algorithm used to perform the resize. Choices range from simple interpolation (Lanczos, Nearest) to upscaler models (ESRGAN, SwinIR, RealESRGAN, etc.) and latent-space methods. Upscaler models give better quality at the cost of speed; simple methods are fast but soft.","ui":"resize"},
{"id":"","label":"Model repo","localized":"","hint":"HuggingFace repository ID for the model","ui":"script_prompt_enhance"},
{"id":"","label":"Model gguf","localized":"","hint":"Optional GGUF quantized model repository on HuggingFace","ui":"script_prompt_enhance"},
{"id":"","label":"Model type","localized":"","hint":"Optional GGUF model quantization type","ui":"script_prompt_enhance"},
@@ -916,10 +917,10 @@
{"id":"","label":"Min guidance","localized":"","hint":"","ui":"script_video"},
{"id":"","label":"Max guidance","localized":"","hint":"","ui":"script_video"},
{"id":"","label":"Motion level","localized":"","hint":"","ui":"script_video"},
- {"id":"","label":"Mode after","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Method after","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Mode mask","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Method mask","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Mode after","localized":"","hint":"How the output image is fitted to the target resolution after the model finishes generating (Post sub-tab in the Size accordion).None : skip resize, pass the image through unchanged.Fixed : force to target width and height, distorting aspect ratio if they differ.Crop : scale to fully cover the target then center-crop the overflow, preserving aspect ratio.Fill : scale to fit inside the target then pad the remaining space with the background color (set in Settings → Image options).Outpaint : like Fill, but the model paints new content into the padded space instead of using a solid color.Context aware : use seam-carving to add or remove pixels along smooth, featureless paths through the image (like sky or plain backgrounds), preserving the detailed regions. Behavior is controlled by the Context dropdown next to this one.","ui":"control"},
+ {"id":"","label":"Method after","localized":"","hint":"Algorithm used to resize the output image after the model finishes generating (Post sub-tab in the Size accordion). Choices range from simple interpolation (Lanczos, Nearest) to upscaler models (ESRGAN, SwinIR, RealESRGAN, etc.) and latent-space methods. Upscaler models give better quality at the cost of speed; simple methods are fast but soft.","ui":"control"},
+ {"id":"","label":"Mode mask","localized":"","hint":"How the input mask image (used for inpainting, outpainting, or control masks) is fitted to the target resolution (Mask sub-tab in the Size accordion).None : skip resize, pass the mask through unchanged.Fixed : force to target width and height, distorting aspect ratio if they differ.Crop : scale to fully cover the target then center-crop the overflow, preserving aspect ratio.Fill : scale to fit inside the target then pad the remaining space with the background color (set in Settings → Image options).Outpaint : like Fill, but the model paints new content into the padded space instead of using a solid color.Context aware : use seam-carving to add or remove pixels along smooth, featureless paths through the image (like sky or plain backgrounds), preserving the detailed regions. Behavior is controlled by the Context dropdown next to this one.","ui":"control"},
+ {"id":"","label":"Method mask","localized":"","hint":"Algorithm used to resize the input mask image (used for inpainting, outpainting, or control masks; Mask sub-tab in the Size accordion). Choices range from simple interpolation (Lanczos, Nearest) to upscaler models (ESRGAN, SwinIR, RealESRGAN, etc.) and latent-space methods. Upscaler models give better quality at the cost of speed; simple methods are fast but soft.","ui":"control"},
{"id":"","label":"Maximum units","localized":"","hint":"","ui":"control"},
{"id":"","label":"Max faces","localized":"","hint":"","ui":"control"},
{"id":"","label":"Medium","localized":"","hint":"","ui":"control"},
@@ -992,8 +993,8 @@
{"id":"","label":"Noise scale","localized":"","hint":"","ui":"video"},
{"id":"","label":"Note","localized":"","hint":"","ui":"component-8823"},
{"id":"","label":"Non-blocking move operations","localized":"","hint":"","ui":"settings_offload"},
- {"id":"","label":"Nunchaku attention","localized":"","hint":"Replaces default attention with Nunchaku's custom FP16 attention kernel for faster inference on consumer NVIDIA GPUs. Might provide performance improvement on GPUs which have higher FP16 tensor cores throughput than BF16. Currently only affects Flux-based models (Dev, Schnell, Kontext, Fill, Depth, etc.). Has no effect on Qwen, SDXL, Sana, or other architectures. Disabled by default.","ui":"settings_quantization"},
- {"id":"","label":"Nunchaku offloading","localized":"","hint":"Enables Nunchaku's own per-block CPU offloading with asynchronous CUDA streams to reduce VRAM usage. Uses a ping-pong buffer strategy: while one transformer block computes on GPU, the next block preloads from CPU in the background, hiding most of the transfer latency. Can reduce VRAM usage at the cost of slower inference. This replaces SD.Next's pipeline offloading for the transformer component. Only useful on low-VRAM GPUs. If your GPU has enough memory to hold the quantized model (16+ GB), keep this disabled for maximum speed. Supports Flux and Qwen models. Not supported for SDXL where this setting is ignored. Disabled by default.","ui":"settings_quantization"},
+ {"id":"","label":"Nunchaku attention","localized":"","hint":"Replaces default attention with Nunchaku's custom FP16 attention kernel for faster inference on consumer NVIDIA GPUs. Might provide performance improvement on GPUs which have higher FP16 tensor cores throughput than BF16. Currently only affects Flux -based models (Dev , Schnell , Kontext , Fill , Depth , etc.). Has no effect on Qwen , SDXL , Sana , or other architectures. Disabled by default.","ui":"settings_quantization"},
+ {"id":"","label":"Nunchaku offloading","localized":"","hint":"Enables Nunchaku's own per-block CPU offloading with asynchronous CUDA streams to reduce VRAM usage. Uses a ping-pong buffer strategy: while one transformer block computes on GPU, the next block preloads from CPU in the background, hiding most of the transfer latency. Can reduce VRAM usage at the cost of slower inference. This replaces SD.Next's pipeline offloading for the transformer component. Only useful on low-VRAM GPUs. If your GPU has enough memory to hold the quantized model (16+ GB), keep this disabled for maximum speed. Supports Flux and Qwen models. Not supported for SDXL where this setting is ignored. Disabled by default.","ui":"settings_quantization"},
{"id":"","label":"native","localized":"","hint":"","ui":"settings_text_encoder"},
{"id":"","label":"no-grad","localized":"","hint":"Disables gradient tracking with torch.no_grad. Reduces memory usage and speeds up inference.","ui":"settings_backends"},
{"id":"","label":"Numbered filenames","localized":"","hint":"","ui":"settings_saving-paths"},
@@ -1006,7 +1007,7 @@
{"id":"","label":"Network parameters","localized":"","hint":""}
],
"o": [
- {"id":"txt2img_results_mobile","label":"Output","localized":"","hint":"Show/hide selection of output media: generation resuls and live previews during generation process","ui":"txt2img"},
+ {"id":"txt2img_results_mobile","label":"Output","localized":"","hint":"Generation resuls and live previews during generation process Click to minimize/maximize","ui":"txt2img"},
{"id":"","label":"OpenCLiP","localized":"","hint":"Analyze image using CLiP model via OpenCLiP","ui":"caption"},
{"id":"","label":"ONNX","localized":"","hint":""},
{"id":"","label":"Override","localized":"","hint":"Override settings that can change server behavior and are typically applied from imported image metadata","ui":"txt2img"},
@@ -1052,7 +1053,7 @@
{"id":"","label":"olive-ai","localized":"","hint":"","ui":"settings_compile"},
{"id":"","label":"openvino_fx","localized":"","hint":"","ui":"settings_compile"},
{"id":"","label":"Overwrite existing","localized":"","hint":"","ui":"models_current_tab"},
- {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)","ui":"component-5674"},
+ {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5 , 9 values for SDXL )","ui":"component-5674"},
{"id":"","label":"Overwrite model","localized":"","hint":"","ui":"models_merge_tab"},
{"id":"","label":"Output model","localized":"","hint":"","ui":"models_replace_tab"},
{"id":"","label":"Overwrite existing file","localized":"","hint":"","ui":"component-5851"},
@@ -1064,7 +1065,7 @@
{"id":"txt2img_prompts","label":"Prompts","localized":"","hint":"Image prompt and negative prompt","ui":"txt2img"},
{"id":"txt2img_pause","label":"Pause","localized":"","hint":"Pause processing","ui":"txt2img"},
{"id":"","label":"Post","localized":"","hint":"Resize image after processing","ui":"control"},
- {"id":"","label":"Preview","localized":"","hint":"","ui":"video"},
+ {"id":"","label":"Preview","localized":"","hint":"Selects how the mask preview is rendered when you click Run Preview .None : skip the preview step.Masked : input image with everything outside the mask blacked out.Binary : pure black-and-white mask (Otsu thresholded).Grayscale : mask intensity values rendered as gray levels.Color : mask recolored using the selected Colormap .Composite : 50/50 blend of the input image and the colored mask, so you can see exactly where the mask falls relative to the subject. Default Composite.","ui":"video"},
{"id":"","label":"Process Image","localized":"","hint":"Process single image","ui":"extras"},
{"id":"","label":"Process Batch","localized":"","hint":"Process batch of images","ui":"extras"},
{"id":"","label":"Process Folder","localized":"","hint":"Process all images in a folder","ui":"extras"},
@@ -1074,13 +1075,12 @@
{"id":"","label":"Preset Block Merge","localized":"","hint":"","ui":"models_merge_tab"},
{"id":"","label":"Preview metadata","localized":"","hint":""},
{"id":"","label":"Prompt","localized":"","hint":"Describe image you want to generate","ui":"txt2img"},
- {"id":"","label":"Processed Preview","localized":"","hint":"Show/hide section from pre-processing of input images before actual generate","ui":"control"},
{"id":"","label":"PixelArt","localized":"","hint":"","ui":"extras"},
{"id":"","label":"PAG: Perturbed attention guidance","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"PAB: Pyramid attention broadcast","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"Para-attention","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"Paths for specific models","localized":"","hint":"","ui":"settings_system-paths"},
- {"id":"","label":"Prediction method","localized":"","hint":"Defines what the model predicts at each step. Options: - default: the model default - epsilon: noise (most common for Stable Diffusion) - sample: direct denoised image prediction, also called as x0 prediction - v_prediction: velocity prediction, used by CosXL and NoobAI VPred models - flow_prediction: used with newer flow-matching models like SD3 and Flux","ui":"txt2img"},
+ {"id":"","label":"Prediction method","localized":"","hint":"Defines what the model predicts at each step. Options: - default : the model default - epsilon : noise (most common for Stable Diffusion) - sample : direct denoised image prediction, also called as x0 prediction - v_prediction : velocity prediction, used by CosXL and NoobAI VPred models - flow_prediction : used with newer flow-matching models like SD3 and Flux ","ui":"txt2img"},
{"id":"","label":"PAG scale","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"PAG start","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"PAG stop","localized":"","hint":"","ui":"txt2img"},
@@ -1104,10 +1104,9 @@
{"id":"","label":"Preview start","localized":"","hint":"","ui":"script_instantir"},
{"id":"","label":"Preview end","localized":"","hint":"","ui":"script_instantir"},
{"id":"","label":"Pixels to expand","localized":"","hint":"","ui":"script_outpainting"},
- {"id":"","label":"Processor","localized":"","hint":"Processor type to use to preprocess image used for ControlNet","ui":"control"},
+ {"id":"","label":"Processor","localized":"","hint":"Processor type to use to preprocess image used for ControlNet ","ui":"control"},
{"id":"","label":"Pose confidence","localized":"","hint":"","ui":"control"},
{"id":"","label":"Parameter free","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Processed","localized":"","hint":"Show/hide section with processed images","ui":"control"},
{"id":"","label":"Postprocess mask","localized":"","hint":"","ui":"extras"},
{"id":"","label":"PixelArt block size","localized":"","hint":"","ui":"extras"},
{"id":"","label":"PixelArt sharpen","localized":"","hint":"","ui":"extras"},
@@ -1169,7 +1168,7 @@
{"id":"","label":"Resize to","localized":"","hint":"","ui":"control"},
{"id":"","label":"Resize\n by","localized":"","hint":"","ui":"control"},
{"id":"","label":"Resize\n to","localized":"","hint":"","ui":"control"},
- {"id":"control_mask_refresh","label":"Run Preview","localized":"","hint":"","ui":"control"},
+ {"id":"control_mask_refresh","label":"Run Preview","localized":"","hint":"Runs the configured mask pipeline (auto-segment, dilate, erode, blur, invert) on the current input and renders the result in the output panel using the selected Preview style. Use this to iterate on mask settings without launching a full generation.","ui":"control"},
{"id":"","label":"Reference","localized":"","hint":"List of reference models that can be automatically downloaded on first use","ui":"control"},
{"id":"framepack_btn_reset_model","label":"Reset receipe","localized":"","hint":"","ui":"video"},
{"id":"video_generation_info_button","label":"Run","localized":"","hint":"","ui":"video"},
@@ -1194,11 +1193,11 @@
{"id":"","label":"Resize","localized":"","hint":"Image resizing, can be using fixed resolution on based on scale","ui":"settings_postprocessing"},
{"id":"","label":"Rerefence models","localized":"","hint":"","ui":"settings_extra_networks"},
{"id":"","label":"Replace model components","localized":"","hint":"","ui":"models_replace_tab"},
- {"id":"","label":"rescale","localized":"","hint":"rescale betas with zero terminal snr","ui":"txt2img"},
+ {"id":"","label":"rescale","localized":"","hint":"Rescales the noise schedule so the final timestep starts from true pure noise (zero signal-to-noise ratio). Standard SD schedules don't quite reach pure noise at the highest timestep, which biases generations toward medium brightness and limits dynamic range. Rescaling unlocks the full range of darks and brights. Should only be enabled for models trained with zero-terminal-SNR or v-prediction. The most common are SDXL fine-tunes carrying a 'vpred' or 'v-prediction' tag in the name (e.g. NoobAI XL Vpred , Illustrious XL Vpred , Stability's CosXL ), plus some noise-offset and Terminus -family checkpoints. Standard epsilon-prediction models such as base SDXL , Pony , and Animagine should be left as-is. Enabling on a mismatched model will shift colors and degrade quality. Recommended to leave off unless your model documentation specifically calls for it. Disabled by default.","ui":"txt2img"},
{"id":"","label":"Resize seed from width","localized":"","hint":"Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution","ui":"txt2img"},
{"id":"","label":"Resize seed from height","localized":"","hint":"Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution","ui":"txt2img"},
- {"id":"","label":"Refine guidance","localized":"","hint":"CFG scale used for refiner pass","ui":"txt2img"},
- {"id":"","label":"Resize mode","localized":"","hint":"Defines how the input is resized or adapted in second-pass refinement: - none: no resizing, keep original resolution - fixed: force resize to target resolution (may distort) - crop: center-crop to fit target while keeping aspect ratio - fill: resize to fit and pad empty space with borders - outpaint: extend canvas beyond image borders - context aware: smart resize that blends or adapts surrounding areas","ui":"txt2img"},
+ {"id":"","label":"Refine guidance","localized":"","hint":"Guidance scale used for the secondary pass (refiner model or HiRes refine). Behaves like the main Guidance scale but applies only to that secondary pass. For OmniGen this slider controls a separate image-conditioning guidance scale instead, used alongside the main Guidance scale in OmniGen's dual-CFG formula. Set to 0 to disable guidance for the secondary pass. Defaults to 6.0.","ui":"txt2img"},
+ {"id":"","label":"Resize mode","localized":"","hint":"Defines how the input is resized or adapted in second-pass refinement: - none : no resizing, keep original resolution - fixed : force resize to target resolution (may distort) - crop : center-crop to fit target while keeping aspect ratio - fill : resize to fit and pad empty space with borders - outpaint : extend canvas beyond image borders - context aware : smart resize that blends or adapts surrounding areas","ui":"txt2img"},
{"id":"","label":"Resize method","localized":"","hint":"Method used to resize the image: can be simple resize, upscaling model, latent resize or asymmetric decode","ui":"txt2img"},
{"id":"","label":"Resize width","localized":"","hint":"Resizes image to this width. If 0, width is inferred from either of two nearby sliders","ui":"txt2img"},
{"id":"","label":"Resize height","localized":"","hint":"Resizes image to this height. If 0, height is inferred from either of two nearby sliders","ui":"txt2img"},
@@ -1208,8 +1207,8 @@
{"id":"","label":"Refiner steps","localized":"","hint":"Number of steps to use for refiner pass","ui":"txt2img"},
{"id":"","label":"Refine prompt","localized":"","hint":"Prompt used for both second encoder in base model (if it exists) and for refiner pass (if enabled)","ui":"txt2img"},
{"id":"","label":"Refine negative prompt","localized":"","hint":"Negative prompt used for both second encoder in base model (if it exists) and for refiner pass (if enabled)","ui":"txt2img"},
- {"id":"","label":"Renoise","localized":"","hint":"Apply additional noise during detailing","ui":"txt2img"},
- {"id":"","label":"Renoise end","localized":"","hint":"Final step when renoise is applied","ui":"txt2img"},
+ {"id":"","label":"Renoise","localized":"","hint":"Multiplier applied to the sampler's step size during the detailer pass. Same mechanism as the Sigma adjust slider in the sampler tab, scoped to detailer only. Values below 1.0 shrink each step for smoother, more conservative refinement (good for keeping faces stable). Values above 1.0 enlarge each step for sharper, more aggressive resampling. Default 1.0 disables the adjustment.","ui":"txt2img"},
+ {"id":"","label":"Renoise end","localized":"","hint":"Upper bound of the denoising window where Renoise is active within the detailer pass, as a fraction of the noise schedule (1.0 = pure noise, 0.0 = clean image). Lower values restrict renoise to the very first steps (gentler intervention); higher values let it act further into the run. Default 1.0 keeps renoise active across the full pass.","ui":"txt2img"},
{"id":"","label":"Repeat x-axis","localized":"","hint":"","ui":"script_asymmetric_tiling"},
{"id":"","label":"Repeat y-axis","localized":"","hint":"","ui":"script_asymmetric_tiling"},
{"id":"","label":"ReSwapper Model","localized":"","hint":"","ui":"script_face"},
@@ -1237,7 +1236,6 @@
{"id":"","label":"RAS enabled","localized":"","hint":"","ui":"settings_advanced"},
{"id":"","label":"reduce-overhead","localized":"","hint":"","ui":"settings_compile"},
{"id":"","label":"repeated","localized":"","hint":"","ui":"settings_compile"},
- {"id":"","label":"Replace underscores","localized":"","hint":"Display underscores in tag names as spaces in the autocomplete suggestion list. For example, long_hair appears as long hair .","ui":"script_autocomplete"},
{"id":"","label":"Root model folder","localized":"","hint":"","ui":"settings_system-paths"},
{"id":"","label":"Resize background color","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"Restore from metadata: skip params","localized":"","hint":"","ui":"settings_image-metadata"},
@@ -1302,9 +1300,9 @@
{"id":"","label":"Server log","localized":"","hint":""},
{"id":"","label":"Steps","localized":"","hint":"How many times to improve the generated image iteratively; higher values take longer; very low values can produce bad results","ui":"txt2img"},
{"id":"","label":"Sampling method","localized":"","hint":"Which algorithm to use to produce the image","ui":"txt2img"},
- {"id":"","label":"Sigma method","localized":"","hint":"Controls how noise levels (sigmas) are distributed across diffusion steps. Options: - default: the model default - karras: smoother noise schedule, higher quality with fewer steps - beta: based on beta schedule values - exponential: exponential decay of noise - lambdas: experimental, balances signal-to-noise - flowmatch: tuned for flow-matching models","ui":"txt2img"},
- {"id":"","label":"Sigma adjust","localized":"","hint":"Adjust sampler sigma value","ui":"txt2img"},
- {"id":"","label":"Sampler order","localized":"","hint":"Order of solver updates in the sampler. Higher order improves stability/accuracy but increases compute cost.","ui":"txt2img"},
+ {"id":"","label":"Sigma method","localized":"","hint":"Controls how noise levels (sigmas) are distributed across diffusion steps.Default : use the scheduler's built-in sigma method.Karras : smoother schedule that emphasizes later steps where fine details emerge; generally higher quality with fewer steps.Betas : derive sigmas directly from the model's beta schedule (classic DDPM behavior).Exponential : exponential decay of noise across steps; aggressive denoising early, slower refinement later.Lambdas : Lu's lambdas method from the DPM-Solver paper, specific to the DPM++ family.Flowmatch : sigma schedule tuned for flow-matching models (Flux , SD3 , video models).","ui":"txt2img"},
+ {"id":"","label":"Sigma adjust","localized":"","hint":"Multiplier applied to the sampler's step size during the active timestep window. (Sigma is the amount of noise the sampler removes at each step.) Values below 1.0 shrink the step for smoother, more conservative denoising. Values above 1.0 enlarge it for sharper, more aggressive sampling. Default 1.0 disables the adjustment entirely. Use Adjust start and Adjust end to define the timestep range where the multiplier takes effect.","ui":"txt2img"},
+ {"id":"","label":"Sampler order","localized":"","hint":"Overrides the solver order of the active sampler when set above 0. Higher orders use more historical steps per update for greater stability and accuracy at the cost of extra compute. Lower orders are faster but noisier. Default 0 leaves each sampler at its built-in order. Many samplers in the dropdown already encode their order in the name (e.g. DPM++ 2M is order 2, DPM++ 3M is order 3, DPM++ 2M SDE is order 2). Within a sampler family, the named variants differ ONLY by this value, so picking DPM++ 2M with the slider at 3 produces a scheduler that is functionally identical to picking DPM++ 3M with the slider at 0. The same equivalence holds across the rest of the DPM++ multistep family (including the SDE and Inverse variants) and across the ER-SDE family. Samplers without a configurable solver order (DDIM , plain Euler , ancestrals, etc.) ignore this slider entirely.","ui":"txt2img"},
{"id":"","label":"SLG scale","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"SLG start","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"SLG stop","localized":"","hint":"","ui":"txt2img"},
@@ -1318,7 +1316,7 @@
{"id":"","label":"SEG layers","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"SEG config","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Strength","localized":"","hint":"Denoising strength of during image operation controls how much of original image is allowed to change during generate","ui":"txt2img"},
- {"id":"","label":"Sort detections","localized":"","hint":"Sort detected areas by from left to right instead of detection score","ui":"txt2img"},
+ {"id":"","label":"Sort detections","localized":"","hint":"Process detected regions left-to-right (by bounding box X position) instead of in detection-score order. Improves consistency when the prompt assigns different traits to different subjects in a multi-line prompt: prompts are mapped per detection in order, so a stable left-to-right order makes line 1 always go to the leftmost subject. Default off.","ui":"txt2img"},
{"id":"","label":"Saturation","localized":"","hint":"Controls color intensity. Positive values make colors more vivid, negative values desaturate toward grayscale. At -1.0 the image becomes fully monochrome.","ui":"txt2img"},
{"id":"","label":"Sharpness","localized":"","hint":"Enhances edge detail and fine textures. Higher values produce crisper edges but may amplify noise or artifacts if pushed too far. Set to 0 to disable. Operates via an unsharp mask kernel.","ui":"txt2img"},
{"id":"","label":"Shadows","localized":"","hint":"Adjusts the brightness of shadow (dark) regions. Positive values lift shadows to reveal detail, negative values deepen them. Operates on the L channel in Lab color space using a luminance-weighted mask, leaving highlights and midtones largely unaffected.","ui":"txt2img"},
@@ -1350,7 +1348,8 @@
{"id":"","label":"Show input","localized":"","hint":"","ui":"control"},
{"id":"","label":"Show preview","localized":"","hint":"","ui":"control"},
{"id":"","label":"Separate init image","localized":"","hint":"Creates an additional window next to Control input labeled Init input, so you can have a separate image for both Control operations and an init source.","ui":"control"},
- {"id":"","label":"Skip input frames","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Skip input processing","localized":"","hint":"Bypasses the active control processor and feeds the raw input image directly to the pipeline. Use when you have already preprocessed the image externally (depth map, canny edges, openpose skeleton, etc.) and don't want SD.Next to re-run the processor on it. The input still routes through any selected ControlNet /T2I-Adapter /etc. model, just without the preprocessing step. Default off.","ui":"control"},
+ {"id":"","label":"Skip input frames","localized":"","hint":"Number of input frames to skip between each processed frame when the input is a video. Use to thin out long source videos: only every (N+1)-th frame is processed and the rest are dropped. Set to 0 to process every frame. Set to 1 to process every other frame, 2 for every third, and so on. Default 0.","ui":"control"},
{"id":"","label":"Style fidelity","localized":"","hint":"","ui":"control"},
{"id":"","label":"Scribble","localized":"","hint":"","ui":"control"},
{"id":"","label":"Score threshold","localized":"","hint":"","ui":"control"},
@@ -1426,7 +1425,7 @@
{"id":"","label":"SDXL","localized":"","hint":"StableDiffusion XL","ui":"component-5660"},
{"id":"","label":"Save metadata","localized":"","hint":"","ui":"models_merge_tab"},
{"id":"","label":"safetensors","localized":"","hint":"","ui":"models_merge_tab"},
- {"id":"","label":"shuffle","localized":"","hint":"Loads full model in RAM and calculates on VRAM: Less speedup, suggested for SDXL merges","ui":"models_merge_tab"},
+ {"id":"","label":"shuffle","localized":"","hint":"Loads full model in RAM and calculates on VRAM: Less speedup, suggested for SDXL merges","ui":"models_merge_tab"},
{"id":"","label":"Save diffusers","localized":"","hint":"","ui":"models_replace_tab"},
{"id":"","label":"Save safetensors","localized":"","hint":"","ui":"models_replace_tab"},
{"id":"","label":"Sort","localized":"","hint":"","ui":"models_civitai_tab"},
@@ -1456,10 +1455,10 @@
{"id":"","label":"Theme options","localized":"","hint":"","ui":"settings_ui"},
{"id":"","label":"Task History","localized":"","hint":""},
{"id":"","label":"Tone","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"Timestep spacing","localized":"","hint":"Determines how timesteps are spaced across the diffusion process. Options: - default: the model default - leading: creates evenly spaced steps - linspace: includes the first and last steps and evenly selects the remaining intermediate steps - trailing: only includes the last step and evenly selects the remaining intermediate steps starting from the end","ui":"txt2img"},
- {"id":"","label":"Timesteps presets","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"Timesteps override","localized":"","hint":"","ui":"txt2img"},
- {"id":"","label":"thresholding","localized":"","hint":"","ui":"txt2img"},
+ {"id":"","label":"Timestep spacing","localized":"","hint":"Determines how timesteps are spaced across the diffusion process. Options: - default : the model default - leading : creates evenly spaced steps - linspace : includes the first and last steps and evenly selects the remaining intermediate steps - trailing : only includes the last step and evenly selects the remaining intermediate steps starting from the end","ui":"txt2img"},
+ {"id":"","label":"Timesteps presets","localized":"","hint":"Picks a hand-tuned timestep schedule and writes it into Timesteps override. 'AYS SD15' and 'AYS SDXL' load the Align Your Steps schedules optimized for those base models, both 10 steps long. Selecting one of these effectively forces the generation to run at exactly 10 steps regardless of the main Steps slider, which is intended: AYS produces results comparable to 30+ step traditional sampling at this length. Use the SD15 preset for SD 1.x checkpoints and the SDXL preset for SDXL -based checkpoints. The AYS schedules are not appropriate for flow-matching models (Flux , SD3 ) or other architectures. Set to None to clear the override. No preset by default.","ui":"txt2img"},
+ {"id":"","label":"Timesteps override","localized":"","hint":"Comma- or space-separated list of integer timesteps in the 0-999 range, listed from highest (most noisy) to lowest (cleanest). When set, this list completely replaces the scheduler's normal timestep schedule and forces the step count to match the list length, ignoring the main Steps slider. Requires at least 3 values to take effect; shorter inputs are silently ignored. Not all samplers support arbitrary timestep injection. If the active sampler doesn't, a warning is logged and the override is skipped. Selecting a preset from Timesteps presets fills this field automatically. Useful for advanced users experimenting with custom schedules. Most users should leave this blank. Clear the field to disable. Empty by default.","ui":"txt2img"},
+ {"id":"","label":"thresholding","localized":"","hint":"Enables dynamic thresholding. At each step the predicted clean image is clipped so its values stay within the model's trained range, which suppresses saturation and washed-out colors at high guidance. Most useful for SD 1.x and SD 2.x at high CFG (>10). Generally not helpful for SDXL or flow-matching models, which already handle high CFG gracefully. Applies to DPM++ family, UniPC , DDIM , DEIS , SA Solver , and DC Solver . Recommended to leave off unless you see saturation artifacts. Disabled by default.","ui":"txt2img"},
{"id":"","label":"Tint strength","localized":"","hint":"","ui":"txt2img"},
{"id":"","label":"Texture tiling","localized":"","hint":"Apply seamless tiling to generated image so it can be used as a texture","ui":"txt2img"},
{"id":"","label":"Threshold","localized":"","hint":"","ui":"script_apg"},
@@ -1496,7 +1495,7 @@
{"id":"","label":"Tiny","localized":"","hint":"","ui":"control"},
{"id":"","label":"True guidance","localized":"","hint":"","ui":"video"},
{"id":"","label":"Tile frames","localized":"","hint":"","ui":"video"},
- {"id":"","label":"Task","localized":"","hint":"Changes which task the model will perform. Regular text prompts can be used when the task is set to Use Prompt . When other options are selected, see the hint text inside an empty Prompt field for guidance.","ui":"caption"},
+ {"id":"","label":"Task","localized":"","hint":"Changes which task the model will perform. Regular text prompts can be used when the task is set to Use Prompt . When other options are selected, see the hint text inside an empty Prompt field for guidance.","ui":"caption"},
{"id":"","label":"Tagger Model","localized":"","hint":"Model to use for image tagging.WaifuDiffusion models (wd-*): Modern taggers with separate general and character thresholds.DeepBooru : Legacy tagger, uses only general threshold.","ui":"caption"},
{"id":"","label":"Torch","localized":"","hint":"","ui":"component-8779"},
{"id":"","label":"Transformers load using Run:ai streamer","localized":"","hint":"","ui":"settings_sd"},
@@ -1537,7 +1536,8 @@
{"id":"","label":"Upscale","localized":"","hint":"Upscale image","ui":"extras"},
{"id":"","label":"UI Tabs","localized":"","hint":"","ui":"settings_ui"},
{"id":"","label":"Upscaling","localized":"","hint":"","ui":"settings_postprocessing"},
- {"id":"","label":"Use segmentation","localized":"","hint":"Run detailer using segmentation mask","ui":"txt2img"},
+ {"id":"","label":"Use segmentation","localized":"","hint":"Use the model's pixel-precise segmentation mask as the inpaint mask instead of the rectangular bounding box. Tighter mask means less unintended change around the detection (e.g., the inpaint stays on the face, not on the hair or background behind it). Better blending and smaller seams. Requires a segmentation-capable model (filename usually contains -seg). Bounding-box-only models silently fall back to the rectangle. Default off.","ui":"txt2img"},
+ {"id":"","label":"Use init image","localized":"","hint":"Decides whether the input image is also used as an init image for img2img-style modification.No: Control only : the input is used only by the active control processor (depth, canny, pose, etc.) to guide the model; the picture itself is built from scratch by the model. Standard ControlNet behavior.1st: Same as control : the control input doubles as the init image, so the model starts from your image and modifies it instead of building one from scratch. Useful for inpainting, restyling, or adding control guidance to img2img with a single source image.2nd: Separate image : opens an extra Init input pane next to Control input so you can supply different sources for control conditioning and img2img init.Denoising strength controls how far the result moves from the init image and only takes effect in the two init modes. Default No: Control only .","ui":"control"},
{"id":"","label":"Unload adapter","localized":"","hint":"Unload IP adapter immediately after generate. Otherwise IP adapter will remain loaded for faster use in next generate process","ui":"txt2img"},
{"id":"","label":"Use same seed","localized":"","hint":"","ui":"script_prompts_from_file"},
{"id":"","label":"Use defaults","localized":"","hint":"","ui":"script_video"},
@@ -1600,7 +1600,7 @@
{"id":"","label":"Vignette","localized":"","hint":"Applies radial edge darkening that draws focus toward the center of the image. Higher values produce a stronger falloff from center to corners. Set to 0 to disable. Simulates the natural light falloff seen in vintage and cinematic lenses.","ui":"txt2img"},
{"id":"","label":"VAE type","localized":"","hint":"Choose if you want to run full VAE, reduced quality VAE or attempt to use remote VAE service","ui":"txt2img"},
{"id":"","label":"Version","localized":"","hint":"","ui":"script_pulid"},
- {"id":"","label":"Video format","localized":"","hint":"Format and codec of output video","ui":"script_video"},
+ {"id":"","label":"Video format","localized":"","hint":"Container format and codec for the output video file. Pick a format your downstream tools understand. MP4/MP4V is broadly compatible with most players and editors. Other choices trade off file size, quality, and player support. Default MP4/MP4V.","ui":"script_video"},
{"id":"","label":"Video duration","localized":"","hint":"","ui":"script_video"},
{"id":"","label":"Video engine","localized":"","hint":"","ui":"video"},
{"id":"","label":"Video model","localized":"","hint":"","ui":"video"},
@@ -1632,10 +1632,10 @@
{"id":"","label":"Wildcards","localized":"","hint":""},
{"id":"","label":"WanAI","localized":"","hint":"","ui":"settings_model_options"},
{"id":"","label":"Watermarking","localized":"","hint":"","ui":"settings_saving-images"},
- {"id":"","label":"Width","localized":"","hint":"Image width","ui":"txt2img"},
+ {"id":"","label":"Width","localized":"","hint":"Target width of the output image in pixels. For generation, this sets the resolution the model produces. For resize and upscale operations, this is the width the input is fitted to. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"txt2img"},
{"id":"","label":"Weight","localized":"","hint":"","ui":"script_resadapter"},
- {"id":"","label":"Width after","localized":"","hint":"","ui":"control"},
- {"id":"","label":"Width mask","localized":"","hint":"","ui":"control"},
+ {"id":"","label":"Width after","localized":"","hint":"Target width of the output image in pixels, applied after the model finishes generating (Post sub-tab in the Size accordion). Use this to upscale or downscale the final image before saving. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"control"},
+ {"id":"","label":"Width mask","localized":"","hint":"Target width of the input mask image in pixels (Mask sub-tab in the Size accordion). The mask is used for inpainting, outpainting, or as a control mask, and is resized so it aligns with the processing resolution. Should be a multiple of 8 for SD1.x and SDXL latents; newer architectures (Flux , SD3 , video models) may require higher multiples (16, 32, or 64). Values that don't match are automatically floored to the nearest valid multiple for the loaded model.","ui":"control"},
{"id":"","label":"WebP lossless compression","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"wavelet","localized":"","hint":"","ui":"settings_postprocessing"},
{"id":"","label":"Weights clip","localized":"","hint":"Forced merged weights to be no heavier than the original model, preventing burn in and overly saturated models","ui":"models_merge_tab"}
diff --git a/html/locale_es.json b/html/locale_es.json
index 42f8f72bb..94d82fbe0 100644
--- a/html/locale_es.json
+++ b/html/locale_es.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Medios de entrada",
"reload": "",
"hint": "Añadir imagen de entrada para ser utilizada para el procesamiento de imagen a imagen, inpaint o control"
diff --git a/html/locale_fr.json b/html/locale_fr.json
index 470792ab9..5b2c37935 100644
--- a/html/locale_fr.json
+++ b/html/locale_fr.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Média d'entrée",
"reload": "",
"hint": "Ajouter une image d'entrée à utiliser pour le traitement image-à-image, inpaint ou control"
diff --git a/html/locale_he.json b/html/locale_he.json
index 8d1e7704f..27a37076b 100644
--- a/html/locale_he.json
+++ b/html/locale_he.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "מדית קלט",
"reload": "",
"hint": "הוספת תמונת קלט לשימוש עבור עיבוד תמונה-לתמונה, מילוי או בקרה"
diff --git a/html/locale_hi.json b/html/locale_hi.json
index 54d213b59..8aa989b29 100644
--- a/html/locale_hi.json
+++ b/html/locale_hi.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "इनपुट मीडिया",
"reload": "n/a",
"hint": "इमेज-टू-इमेज, इनपेंट या कंट्रोल प्रोसेसिंग के लिए उपयोग की जाने वाली इनपुट छवि जोड़ें"
diff --git a/html/locale_hr.json b/html/locale_hr.json
index d092a3e55..efa35485f 100644
--- a/html/locale_hr.json
+++ b/html/locale_hr.json
@@ -4824,7 +4824,7 @@
},
{
"id": 0,
- "label": "Input Media",
+ "label": "Input",
"localized": "Ulazni medij",
"reload": "",
"hint": "Dodajte ulaznu sliku koja će se koristiti za image-to-image, inpaint ili kontrolnu obradu"
diff --git a/html/locale_id.json b/html/locale_id.json
index c13c3fe19..7450256f8 100644
--- a/html/locale_id.json
+++ b/html/locale_id.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Media Masukan",
"reload": "n/a",
"hint": "Tambahkan gambar masukan untuk digunakan dalam pemrosesan image-to-image, inpaint, atau kontrol"
diff --git a/html/locale_it.json b/html/locale_it.json
index c557311c1..034654b91 100644
--- a/html/locale_it.json
+++ b/html/locale_it.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "Media di input",
"reload": "n/a",
"hint": "Aggiungi un'immagine di input da utilizzare per elaborazioni image-to-image, inpaint o di controllo"
diff --git a/html/locale_ja.json b/html/locale_ja.json
index d16520c68..e30dba4f0 100644
--- a/html/locale_ja.json
+++ b/html/locale_ja.json
@@ -4817,9 +4817,9 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "入力メディア",
- "reload": "Input Media",
+ "reload": "Input",
"hint": "画像間変換、インペイント、またはコントロール処理に使用する入力画像を追加します"
},
{
diff --git a/html/locale_ko.json b/html/locale_ko.json
index 00e3a4032..a5918ee95 100644
--- a/html/locale_ko.json
+++ b/html/locale_ko.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "입력 미디어",
"reload": "n/a",
"hint": "이미지 대 이미지, 인페인트 또는 제어 처리에 사용할 입력 이미지 추가"
diff --git a/html/locale_nb.json b/html/locale_nb.json
index 1542f19f6..89e39b0f7 100644
--- a/html/locale_nb.json
+++ b/html/locale_nb.json
@@ -4817,8 +4817,8 @@
},
{
"id": 13,
- "label": "Input Media",
- "localized": "Input Media",
+ "label": "Input",
+ "localized": "Input",
"reload": "",
"hint": "Add an image here to use it as a base for editing or guiding the AI."
},
diff --git a/html/locale_po.json b/html/locale_po.json
index 8d8055819..b391db6af 100644
--- a/html/locale_po.json
+++ b/html/locale_po.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Media wejściowe",
"reload": "",
"hint": "Dodaj obraz wejściowy do użycia w przetwarzaniu typu image-to-image, inpaint lub control"
diff --git a/html/locale_pt.json b/html/locale_pt.json
index db261911e..065887bf9 100644
--- a/html/locale_pt.json
+++ b/html/locale_pt.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Mídia de Entrada",
"reload": "",
"hint": "Adicionar imagem de entrada a ser usada para processamento de imagem-para-imagem, inpaint ou controle"
diff --git a/html/locale_qq.json b/html/locale_qq.json
index b033d7be9..1d351b192 100644
--- a/html/locale_qq.json
+++ b/html/locale_qq.json
@@ -4817,7 +4817,7 @@
},
{
"id": 0,
- "label": "Input Media",
+ "label": "Input",
"localized": "Media Input",
"reload": "n/a",
"hint": "Addere imaginem input adhibendam pro processu imaginis-ad-imaginem, inpaint, vel moderationis"
diff --git a/html/locale_ru.json b/html/locale_ru.json
index db706a1a2..318e7754f 100644
--- a/html/locale_ru.json
+++ b/html/locale_ru.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "Входные медиаданные",
"reload": "",
"hint": "Добавить входное изображение для использования в image-to-image, inpaint или для управления генерацией"
diff --git a/html/locale_sr.json b/html/locale_sr.json
index d0df82d7b..2cbccaada 100644
--- a/html/locale_sr.json
+++ b/html/locale_sr.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "Ulazni medij",
"reload": "",
"hint": "Dodajte ulaznu sliku koja će se koristiti za obradu slike-u-sliku, inpaint ili kontrolnu obradu"
diff --git a/html/locale_tb.json b/html/locale_tb.json
index 53f4ad85a..bb09b9850 100644
--- a/html/locale_tb.json
+++ b/html/locale_tb.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Source Telemetry",
"reload": "",
"hint": "Upload and synchronize source data for recursive generation, delta-patching, or neural guidance processing"
diff --git a/html/locale_tlh.json b/html/locale_tlh.json
index 80e758920..a0447e4c8 100644
--- a/html/locale_tlh.json
+++ b/html/locale_tlh.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "nI' Media",
"reload": "",
"hint": "nI' media"
diff --git a/html/locale_tr.json b/html/locale_tr.json
index 275409c4c..2a0ddfd06 100644
--- a/html/locale_tr.json
+++ b/html/locale_tr.json
@@ -4817,7 +4817,7 @@
},
{
"id": 14,
- "label": "Input Media",
+ "label": "Input",
"localized": "Giriş Medyası",
"reload": "",
"hint": "Görüntüden görüntüye, inpaint veya kontrol işleme için kullanılacak giriş görüntüsünü ekleyin"
diff --git a/html/locale_ur.json b/html/locale_ur.json
index 5063f6e7e..4edfe260a 100644
--- a/html/locale_ur.json
+++ b/html/locale_ur.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "ان پٹ میڈیا",
"reload": "",
"hint": "تصویر سے تصویر (image-to-image)، ان پینٹ یا کنٹرول پروسیسنگ کے لیے استعمال ہونے والی ان پٹ تصویر شامل کریں"
diff --git a/html/locale_vi.json b/html/locale_vi.json
index 4e631bb43..06187acbd 100644
--- a/html/locale_vi.json
+++ b/html/locale_vi.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Phương tiện đầu vào",
"reload": "",
"hint": "Thêm hình ảnh đầu vào để sử dụng cho xử lý hình ảnh thành hình ảnh, inpaint hoặc điều khiển"
diff --git a/html/locale_xx.json b/html/locale_xx.json
index 5674c5994..79d813296 100644
--- a/html/locale_xx.json
+++ b/html/locale_xx.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "Eniga Amaskomunikilaro",
"reload": "",
"hint": "Aldoni enigeblan bildon por esti uzata por bild-al-bilda, inpaint aŭ kontrola prilaborado"
diff --git a/html/locale_zh.json b/html/locale_zh.json
index 7aa02402e..700de6f4e 100644
--- a/html/locale_zh.json
+++ b/html/locale_zh.json
@@ -4817,7 +4817,7 @@
},
{
"id": 13,
- "label": "Input Media",
+ "label": "Input",
"localized": "输入媒体",
"reload": "",
"hint": "添加用于图生图、重绘或控制处理的输入图像"
diff --git a/installer.py b/installer.py
index b9cf4ff0c..1f658ae2e 100644
--- a/installer.py
+++ b/installer.py
@@ -37,7 +37,7 @@ log = logging.getLogger('sdnext.installer')
debug = log.debug if os.environ.get('SD_INSTALL_DEBUG', None) is not None else lambda *args, **kwargs: None
setuptools, distutils = None, None # defined via ensure_base_requirements
current_branch = None
-pip_log = '--log pip.log ' if os.environ.get('SD_PIP_DEBUG', None) is not None else ''
+pip_log = '--log pip.log' if os.environ.get('SD_PIP_DEBUG', None) is not None else ''
log_file = os.path.join(os.path.dirname(__file__), 'sdnext.log')
hostname = socket.gethostname()
log_rolled = False
@@ -75,7 +75,7 @@ extensions_commit = { # force specific commit for extensions
'adetailer': 'a89c01d'
# 'stable-diffusion-webui-images-browser': '27fe4a7',
}
-control_extensions = [ # 3rd party extensions marked as safe for control ui
+control_extensions = [ # extensions marked as safe for control ui
'NudeNet',
'IP Adapters',
'Remove background',
@@ -244,26 +244,35 @@ def cleanup_broken_packages():
pass
-def pip(arg: str, ignore: bool = False, quiet: bool = True, uv = True) -> tuple[subprocess.CompletedProcess, str]:
+def pip(arg: str, ignore: bool = False, quiet: bool = True, *, uv = True, constraints = True) -> tuple[subprocess.CompletedProcess | None, str]:
t_start = time.time()
originalArg = arg
- arg = arg.replace('>=', '==')
+ arg = arg.replace('>=', '==').strip()
if opts.get('offline_mode', False):
log.warning('Offline mode enabled')
return None, 'offline'
- package = arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace("--force-reinstall", "").replace(" ", " ").strip()
+ package = arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace("--force-reinstall", "").strip()
uv = uv and args.uv and not package.startswith('git+')
pipCmd = "uv pip" if uv else "pip"
if not quiet and '-r ' not in arg:
log.info(f'Install: package="{package}" mode={"uv" if uv else "pip"}')
- env_args = os.environ.get("PIP_EXTRA_ARGS", "")
- all_args = f'{pip_log}{arg} {env_args}'.strip()
+ env_args = os.environ.get("PIP_EXTRA_ARGS", "").strip()
+ all_args: list[str] = []
+ if pip_log:
+ all_args.append(pip_log)
+ all_args.append(arg)
+ if env_args:
+ all_args.append(env_args)
+ if constraints and "-c " not in env_args:
+ all_args.append("-c constraints.txt")
if not quiet:
- log.debug(f'Running: {pipCmd}="{all_args}"')
- result, output = run(sys.executable, "-m", pipCmd, all_args)
+ log.debug(f'Running: {pipCmd}="{" ".join(all_args)}"')
+
+ result, output = run(sys.executable, "-m", pipCmd, *all_args)
+
if len(result.stderr) > 0:
if uv and result.returncode != 0:
- log.warning(f'Install: cmd="{pipCmd}" args="{all_args}" cannot use uv, fallback to pip')
+ log.warning(f'Install: cmd="{pipCmd}" args="{" ".join(all_args)}" cannot use uv, fallback to pip')
debug(f'Install: uv pip error: {result.stderr}')
cleanup_broken_packages()
return pip(originalArg, ignore, quiet, uv=False)
@@ -485,7 +494,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all:
return
- target_commit = "0f1abc4ae8b0eb2a3b40e82a310507281144c423" # diffusers commit hash == 0.37.1.dev-0427
+ target_commit = "015da50b40ee7a082ea8c17a8c43dff717c9653e" # diffusers commit hash == 0.37.1.dev-0427
# if args.use_rocm or args.use_zluda or args.use_directml:
# sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
pkg = package_spec('diffusers')
@@ -521,7 +530,7 @@ def check_transformers():
else:
# target_transformers = '4.57.6'
target_transformers = None
- target_tokenizers = '0.22.2'
+ target_tokenizers = '0.23.1'
if target_transformers is not None:
# Pinned release version (e.g. DirectML)
if (pkg_transformers is None) or ((pkg_transformers.version != target_transformers) or (pkg_tokenizers is None) or ((pkg_tokenizers.version != target_tokenizers) and (not args.experimental))):
@@ -1279,7 +1288,6 @@ def install_requirements():
# set environment variables controling the behavior of various libraries
def set_environment():
log.debug('Setting environment tuning')
- os.environ.setdefault('PIP_CONSTRAINT', 'constraints.txt')
os.environ.setdefault('ACCELERATE', 'True')
os.environ.setdefault('ATTN_PRECISION', 'fp16')
os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100')
@@ -1374,7 +1382,16 @@ def get_version(force=False):
try:
origin = run('git', 'remote get-url origin', check=True)[0].stdout
branch_name = run('git', 'rev-parse --abbrev-ref HEAD', check=True)[0].stdout
- version['url'] = origin.removesuffix('.git') + '/tree/' + branch_name
+ # normalize ssh remotes (git@host:owner/repo) and ssh-protocol remotes
+ # (ssh://git@host/owner/repo) to the canonical https form so downstream
+ # url parsers don't have to special-case each remote shape
+ if origin.startswith('git@'):
+ host, _, path = origin.partition(':')
+ origin = f'https://{host[4:]}/{path}'
+ elif origin.startswith('ssh://'):
+ origin = 'https://' + origin[len('ssh://'):].split('@', 1)[-1]
+ origin = origin.removesuffix('.git')
+ version['url'] = origin + '/tree/' + branch_name
version['branch'] = branch_name
if version['branch'] == 'HEAD':
log.warning('Version: detached state detected')
@@ -1530,7 +1547,10 @@ def check_version(reset=True): # pylint: disable=unused-argument
api_base = f'https://api.github.com/repos/{url_parts}'
else:
api_base = 'https://api.github.com/repos/vladmandic/sdnext'
- branches = requests.get(f'{api_base}/branches', timeout=10).json()
+ branches = requests.get(f'{api_base}/branches', timeout=5).json()
+ if not isinstance(branches, list):
+ log.error(f'Repository: branches API returned {branches!r} from {api_base}')
+ return
branch_names = [b['name'] for b in branches if 'name' in b]
log.trace(f'Repository branches: active={branch_name} available={branch_names}')
except Exception as e:
@@ -1541,7 +1561,7 @@ def check_version(reset=True): # pylint: disable=unused-argument
ts('latest', t_start)
return
try:
- commits = requests.get(f'{api_base}/branches/{branch_name}', timeout=10).json()
+ commits = requests.get(f'{api_base}/branches/{branch_name}', timeout=5).json()
latest = commits['commit']['sha']
if len(latest) != 40:
log.error(f'Repository error: commit={latest} invalid')
diff --git a/javascript/autocomplete.js b/javascript/autocomplete.js
index e6e9d2bc7..b020490c2 100644
--- a/javascript/autocomplete.js
+++ b/javascript/autocomplete.js
@@ -358,7 +358,7 @@ function insertExtraNetwork(textarea, item, kind) {
}
/** Insert a tag at the current word position, replacing the typed prefix. */
-function insertTag(textarea, tagName) {
+function insertTag(textarea, tagName, kind = 'tag') {
const info = getCurrentWord(textarea);
if (!info || (info.mode !== 'tag' && info.mode !== 'artist')) return;
const { value } = textarea;
@@ -371,13 +371,13 @@ function insertTag(textarea, tagName) {
const prefix = needsSepBefore ? `${sep} ` : '';
let suffix = `${sep} `;
if (after.length > 0 && after.trimStart().startsWith(',')) suffix = ' ';
- // Artist mode: optionally keep the `@` prefix (Anima syntax); always convert underscores to spaces
- // since Anima requires space-separated artist names. The `@` is consumed for non-Anima models.
+ // Embedding names are file-system identifiers, so underscores must be preserved regardless of the user setting.
+ // Tags and artists honor `autocomplete_keep_underscores`; default is to swap `_` for space.
+ const keepUnderscores = window.opts?.autocomplete_keep_underscores ?? false;
let body = tagName;
- if (info.mode === 'artist') {
- body = body.replace(/_/g, ' ');
- if (window.opts?.autocomplete_at_prefix_artist) body = `@${body}`;
- }
+ if (kind !== 'embed' && !keepUnderscores) body = body.replace(/_/g, ' ');
+ // Artist mode optionally keeps the `@` prefix (Anima syntax). The `@` is consumed for non-Anima models.
+ if (info.mode === 'artist' && window.opts?.autocomplete_at_prefix_artist) body = `@${body}`;
const insertion = `${prefix}${escapeParensForPrompt(body)}${suffix}`;
textarea.value = before.trimEnd() + (before.trimEnd().length > 0 ? ' ' : '') + insertion + after.trimStart();
// Position cursor after the inserted tag + separator
@@ -447,7 +447,7 @@ const dropdown = {
},
render() {
- const replaceUnderscores = window.opts?.autocomplete_replace_underscores ?? true;
+ const keepUnderscores = window.opts?.autocomplete_keep_underscores ?? false;
const queryNorm = this.query.toLowerCase().replace(/ /g, '_');
this.listEl.replaceChildren();
this.results.forEach((tag, i) => {
@@ -462,7 +462,9 @@ const dropdown = {
dot.title = kind === 'tag' ? (engine.categoryNames[tag.category] || '') : kind;
const name = document.createElement('span');
name.className = 'autocomplete-tag';
- const tagText = replaceUnderscores ? tag.display.replace(/_/g, ' ') : tag.display;
+ // Embeddings are file-name identifiers, so they always render as-is to match how they get inserted.
+ const swapForKind = kind !== 'embed';
+ const tagText = (swapForKind && !keepUnderscores) ? tag.display.replace(/_/g, ' ') : tag.display;
const canonicalMatch = tag.name.indexOf(queryNorm);
if (canonicalMatch >= 0 && queryNorm.length > 0) {
const mark = document.createElement('mark');
@@ -480,7 +482,7 @@ const dropdown = {
if (tag.matchedVia === 'alias') annotationTerm = tag.matchedAlias;
else if (tag.matchedVia === 'translation') annotationTerm = tag.matchedTerm;
if (annotationTerm) {
- const annotationDisplay = replaceUnderscores ? annotationTerm.replace(/_/g, ' ') : annotationTerm;
+ const annotationDisplay = (swapForKind && !keepUnderscores) ? annotationTerm.replace(/_/g, ' ') : annotationTerm;
const annotationLower = annotationTerm.toLowerCase();
const annotationMatch = annotationLower.indexOf(queryNorm);
const prefix = tag.matchedVia === 'translation' ? ' \u{1F310} ' : ' (';
@@ -561,7 +563,7 @@ const dropdown = {
insertExtraNetwork(this.textarea, result, result.kind);
} else {
// 'embed' kind and untagged tag results both go through insertTag (comma-aware, paren-escaped).
- insertTag(this.textarea, result.display ?? result.name);
+ insertTag(this.textarea, result.display ?? result.name, result.kind);
}
}
this.hide();
@@ -745,24 +747,24 @@ async function initAutocomplete() {
log('autoComplete', { active, enabled });
// Inject styles (CSS files in javascript/ are not auto-loaded)
const style = document.createElement('style');
- style.textContent = [
- '.autocompleteResults { position: fixed; z-index: 9999; max-height: 300px; overflow-y: auto;',
- ' background: var(--sd-main-background-color, var(--background-fill-primary, #1f2937));',
- ' border: 1px solid var(--sd-input-border-color, var(--border-color-primary, #374151));',
- ' border-radius: var(--sd-border-radius, 6px); box-shadow: 0 4px 16px rgba(0,0,0,0.4);',
- ' font-size: 13px; scrollbar-width: thin; }',
- '.autocompleteResultsList { list-style: none; margin: 0; padding: 4px 0; }',
- '.autocompleteResultsList > li { display: flex; align-items: center; padding: 6px 12px; cursor: pointer;',
- ' gap: 8px; line-height: 1.4; transition: background 0.1s ease; border-bottom: 1px solid rgba(255,255,255,0.03); }',
- '.autocompleteResultsList > li:last-child { border-bottom: none; }',
- '.autocompleteResultsList > li:hover { background: var(--sd-panel-background-color, var(--input-background-fill-focus, #374151)); }',
- '.autocompleteResultsList > li.selected { background: var(--sd-main-accent-color, var(--button-primary-background-fill, #4b5563)); }',
- '.autocomplete-category { font-size: 10px; flex-shrink: 0; width: 10px; text-align: center; cursor: help; }',
- '.autocomplete-tag { flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }',
- '.autocomplete-tag mark { background: transparent; color: inherit; font-weight: 700; }',
- '.autocomplete-count { font-size: 0.75em; opacity: 0.45; flex-shrink: 0; font-variant-numeric: tabular-nums;',
- ' background: rgba(255,255,255,0.06); padding: 1px 6px; border-radius: 8px; min-width: 28px; text-align: right; }',
- ].join('\n');
+ style.textContent = `
+ .autocompleteResults { position: fixed; z-index: 9999; max-height: 300px; overflow-y: auto;
+ background: var(--sd-main-background-color, var(--background-fill-primary, #1f2937));
+ border: 1px solid var(--sd-input-border-color, var(--border-color-primary, #374151));
+ border-radius: var(--sd-border-radius, 6px); box-shadow: 0 4px 16px rgba(0,0,0,0.4);
+ font-size: 13px; scrollbar-width: thin; color: var(--body-text-color-subdued); }
+ .autocompleteResultsList { list-style: none; margin: 0; padding: 4px 0; }
+ .autocompleteResultsList > li { display: flex; align-items: center; padding: 6px 12px; cursor: pointer;
+ gap: 8px; line-height: 1.4; transition: background 0.1s ease; border-bottom: 1px solid rgba(255,255,255,0.03); }
+ .autocompleteResultsList > li:last-child { border-bottom: none; }
+ .autocompleteResultsList > li:hover { background: var(--sd-panel-background-color, var(--input-background-fill-focus, #374151)); }
+ .autocompleteResultsList > li.selected { background: var(--sd-main-accent-color, var(--button-primary-background-fill, #4b5563)); }
+ .autocomplete-category { font-size: 10px; flex-shrink: 0; width: 10px; text-align: center; cursor: help; }
+ .autocomplete-tag { flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+ .autocomplete-tag mark { background: transparent; color: inherit; font-weight: 700; }
+ .autocomplete-count { font-size: 0.75em; opacity: 0.45; flex-shrink: 0; font-variant-numeric: tabular-nums;
+ background: rgba(255,255,255,0.06); padding: 1px 6px; border-radius: 8px; min-width: 28px; text-align: right; }
+ `;
document.head.appendChild(style);
dropdown.init();
await engine.loadEnabled();
diff --git a/javascript/control.js b/javascript/control.js
index 06da73910..1505a6b8d 100644
--- a/javascript/control.js
+++ b/javascript/control.js
@@ -8,11 +8,15 @@ function controlInputMode(inputMode, ...args) {
const tabNames = ['Image', 'Video', 'Batch', 'Folder'];
let inputTab = tabNames[tabIdx] || 'Image';
log('controlInputMode', { mode: inputMode, tab: inputTab, kanvas: typeof Kanvas });
+
+ // if kanvas is available overwrite image inputs with kanvas images
if ((inputTab === 'Image') && (typeof 'Kanvas' !== 'undefined')) {
inputTab = 'Kanvas';
- const imageData = window.kanvas.getImage();
- args[0] = imageData;
+ for (let i = 0; i < window.kanvas.stages.maxStages; i++) {
+ args[4 + i] = window.kanvas.getImage(1 + i, false, false);
+ }
}
+
return [inputTab, ...args];
}
diff --git a/javascript/gallery.js b/javascript/gallery.js
index 377577892..1a123d97a 100644
--- a/javascript/gallery.js
+++ b/javascript/gallery.js
@@ -1,7 +1,11 @@
/* eslint-disable max-classes-per-file */
let ws;
let url;
+let currentSize = 0;
+let currentSort = 'none';
+let currentName = '';
let currentImage = null;
+let currentTitle = '';
let currentGalleryFolder = null;
let pruneImagesTimer;
let outstanding = 0;
@@ -18,7 +22,9 @@ const el = {
search: undefined,
status: undefined,
btnSend: undefined,
+ overlay: undefined,
clearCacheFolder: undefined,
+ size: undefined,
};
const SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'tiff', 'jp2', 'jxl', 'gif', 'mp4', 'mkv', 'avi', 'mjpeg', 'mpg', 'avr'];
@@ -26,12 +32,12 @@ const SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'tiff', 'jp2', 'jxl'
const gallerySorter = {
nameA: { name: 'Name Ascending', func: (a, b) => a.name.localeCompare(b.name) },
nameD: { name: 'Name Descending', func: (b, a) => a.name.localeCompare(b.name) },
- sizeA: { name: 'Size Ascending', func: (a, b) => a.size - b.size },
- sizeD: { name: 'Size Descending', func: (b, a) => a.size - b.size },
- resA: { name: 'Resolution Ascending', func: (a, b) => a.width * a.height - b.width * b.height },
- resD: { name: 'Resolution Descending', func: (b, a) => a.width * a.height - b.width * b.height },
- modA: { name: 'Modified Ascending', func: (a, b) => a.mtime - b.mtime },
- modD: { name: 'Modified Descending', func: (b, a) => a.mtime - b.mtime },
+ sizeD: { name: 'Size Ascending', func: (a, b) => a.size - b.size },
+ sizeA: { name: 'Size Descending', func: (b, a) => a.size - b.size },
+ resD: { name: 'Resolution Ascending', func: (a, b) => a.width * a.height - b.width * b.height },
+ resA: { name: 'Resolution Descending', func: (b, a) => a.width * a.height - b.width * b.height },
+ modD: { name: 'Modified Ascending', func: (a, b) => a.mtime - b.mtime },
+ modA: { name: 'Modified Descending', func: (b, a) => a.mtime - b.mtime },
none: { name: 'None', func: undefined },
};
@@ -71,6 +77,8 @@ function resetGallerySelection() {
updateGallerySelectionClasses(gallerySelection.files, -1);
gallerySelection = { files: [], index: -1 };
currentImage = null;
+ currentName = '';
+ currentTitle = '';
}
function applyGallerySelection(index, { send = true } = {}) {
@@ -84,6 +92,8 @@ function applyGallerySelection(index, { send = true } = {}) {
}
gallerySelection.index = index;
currentImage = files[index].src;
+ currentName = files[index].name;
+ currentTitle = files[index].title;
updateGallerySelectionClasses(files, index);
if (send && el.btnSend) el.btnSend.click();
}
@@ -129,7 +139,7 @@ async function awaitForGallery(expectedSize, signal) {
function updateGalleryStyles() {
if (opts.theme_type?.toLowerCase() === 'modern') {
- folderStylesheet.replaceSync(`
+ folderStylesheet.replace(`
.gallery-folder {
cursor: pointer;
padding: 8px 6px 8px 6px;
@@ -162,7 +172,7 @@ function updateGalleryStyles() {
}
`);
} else {
- folderStylesheet.replaceSync(`
+ folderStylesheet.replace(`
.gallery-folder {
cursor: pointer;
padding: 8px 6px 8px 6px;
@@ -179,16 +189,30 @@ function updateGalleryStyles() {
}
`);
}
- fileStylesheet.replaceSync(`
+ const size = el.size ? el.size.value : opts.extra_networks_card_size;
+ fileStylesheet.replace(`
.gallery-file {
object-fit: contain;
cursor: pointer;
- height: ${opts.extra_networks_card_size}px;
- width: ${opts.browser_fixed_width ? `${opts.extra_networks_card_size}px` : 'unset'};
+ height: ${size}px;
+ width: ${opts.browser_fixed_width ? `${size}px` : 'unset'};
}
.gallery-file:hover {
filter: grayscale(100%);
}
+ .gallery-overlay {
+ position: absolute;
+ height: 24px;
+ background-color: rgba(0,0,0,0.7);
+ display: block;
+ text-align: right;
+ padding: 4px;
+ font-size: 1.2em;
+ letter-spacing: 0.5em;
+ width: 140px;
+ margin-top: calc(140px - 32px);
+ opacity: 75%;
+ }
:host(.gallery-file-selected) .gallery-file {
box-shadow: 0 0 0 2px var(--sd-button-selected-color);
}
@@ -409,9 +433,7 @@ class GalleryFolder extends HTMLElement {
this.div.classList.add('gallery-folder-selected');
GalleryFolder.#active = this;
for (const folder of GalleryFolder.folders) {
- if (folder !== this) {
- folder.div.classList.remove('gallery-folder-selected');
- }
+ if (folder !== this) folder.div.classList.remove('gallery-folder-selected');
}
}
}
@@ -456,7 +478,6 @@ class GalleryFile extends HTMLElement {
this.height = 0;
this.shadow = this.attachShadow({ mode: 'open' });
this.shadow.adoptedStyleSheets = [fileStylesheet];
-
this.firstRun = true;
}
@@ -469,9 +490,7 @@ class GalleryFile extends HTMLElement {
if (dir && dir[1]) {
const dirPath = dir[1];
const isOpen = separatorStates.get(dirPath);
- if (isOpen === false) {
- this.style.display = 'none';
- }
+ if (isOpen === false) this.style.display = 'none';
}
this.hash = await getHash(`${this.src}/${this.size}/${this.mtime}`)
@@ -514,7 +533,7 @@ class GalleryFile extends HTMLElement {
this.size = json.size;
this.mtime = new Date(json.mtime);
if (opts.browser_cache && this.hash) {
- await idbAdd({
+ idbAdd({
hash: this.hash,
folder: this.fullFolder,
file: this.name,
@@ -534,23 +553,30 @@ class GalleryFile extends HTMLElement {
img.src = `file=${this.src}`;
}
}
- if (this.#signal.aborted) { // Do not change the operations order from here...
- return;
- }
+ if (this.#signal.aborted) return;
galleryHashes.add(this.hash);
- if (!ok) {
- return;
- } // ... to here unless modifications are also being made to maintenance functionality and the usage of AbortController/AbortSignal
+ if (!ok) return;
+
img.onclick = () => {
setGallerySelectionByElement(this, { send: true });
};
+ img.onpointerenter = () => {
+ el.overlay.display = 'block';
+ this.shadow.appendChild(el.overlay);
+ currentImage = this.src;
+ currentName = this.name;
+ currentTitle = this.title;
+ };
+ img.onpointerleave = () => {
+ el.overlay.display = 'none';
+ };
img.title = `Folder: ${this.folder}\nFile: ${this.name}\nSize: ${this.size.toLocaleString()} bytes\nModified: ${this.mtime.toLocaleString()}`;
this.title = img.title;
// Final visibility check based on search term.
const shouldDisplayBasedOnSearch = this.title.toLowerCase().includes(el.search.value.toLowerCase());
if (this.style.display !== 'none') { // Only proceed if not already hidden by a closed separator
- this.style.display = shouldDisplayBasedOnSearch ? 'unset' : 'none';
+ this.style.display = shouldDisplayBasedOnSearch ? 'flex' : 'none';
}
this.shadow.appendChild(img);
@@ -558,8 +584,10 @@ class GalleryFile extends HTMLElement {
}
async function createThumb(img) {
- const height = opts.extra_networks_card_size;
- const width = opts.browser_fixed_width ? opts.extra_networks_card_size : 0;
+ const sizeEl = document.getElementById('gallery-thumb-size');
+ currentSize = sizeEl ? parseInt(sizeEl.value, 10) : opts.extra_networks_card_size;
+ const height = currentSize;
+ const width = opts.browser_fixed_width ? currentSize : 0;
const canvas = document.createElement('canvas');
const scaleY = height / img.height;
const scaleX = width > 0 ? width / img.width : scaleY;
@@ -872,8 +900,13 @@ const findDuplicates = (arr, key) => {
};
async function gallerySort(key) {
- if (!Object.hasOwn(gallerySorter, key)) {
- error(`Gallery: "${key}" is not a valid gallery sorting key`);
+ // if currentSort does not start with key, default to key+A
+ // else if currentSort ends with A change to D and vice versa for toggling sort order
+ if (currentSort.startsWith(key)) currentSort = currentSort.endsWith('A') ? `${key}D` : `${key}A`;
+ else currentSort = `${key}A`;
+
+ if (!Object.hasOwn(gallerySorter, currentSort)) {
+ error(`Gallery: "${currentSort}" is not a valid gallery sorting key`);
return;
}
const t0 = performance.now();
@@ -901,7 +934,7 @@ async function gallerySort(key) {
folderGroups.get(dir).push(file);
}
- sortMode = gallerySorter[key];
+ sortMode = gallerySorter[currentSort];
// Sort root files
rootFiles.sort(sortMode.func);
@@ -993,24 +1026,24 @@ async function thumbCacheCleanup(folder, imgCount, controller, force = false) {
if (typeof folder !== 'string' || typeof imgCount !== 'number') {
throw new Error('Function called with invalid arguments');
}
- debug('Thumbnail DB cleanup: Waiting for gallery data to settle');
+ debug('thumbCacheCleanup: wait');
await awaitForGallery(imgCount, controller.signal);
} catch (err) {
- debug(`Thumbnail DB cleanup: Skipping cleanup for "${folder}" due to "${err}"`);
+ error('thumbCacheCleanup', { folder, error: err });
return;
}
maintenanceQueue.enqueue({
signal: controller.signal,
callback: async () => {
- log(`Thumbnail DB cleanup: Checking if "${folder}" needs cleaning`);
+ log('maintenanceQueue', { folder });
const t0 = performance.now();
const keptGalleryHashes = force ? new Set() : new Set(galleryHashes.values()); // External context should be safe since this function run is guarded by AbortController/AbortSignal in the SimpleFunctionQueue
const folderNormalized = folder.replace(/\/+/g, '/').replace(/\/$/, '');
const recursiveFolder = IDBKeyRange.bound(folderNormalized, `${folderNormalized}\uffff`, false, true);
const cachedHashesCount = await idbCount(recursiveFolder)
.catch((e) => {
- error(`Thumbnail DB cleanup: Error when getting entry count for "${folder}".`, e);
+ error('maintenanceQueue', { folder, error: e });
return Infinity; // Forces next check to fail if something went wrong
});
const cleanupCount = cachedHashesCount - keptGalleryHashes.size;
@@ -1020,21 +1053,21 @@ async function thumbCacheCleanup(folder, imgCount, controller, force = false) {
}
if (controller.signal.aborted) {
- debug(`Thumbnail DB cleanup: Cancelling "${folder}" cleanup due to "${controller.signal.reason}"`);
+ debug('maintenanceQueue', { folder, reason: controller.signal.reason });
return;
}
const cb_clearMsg = showCleaningMsg(cleanupCount);
await idbFolderCleanup(keptGalleryHashes, recursiveFolder, controller.signal)
.then((delcount) => {
const t1 = performance.now();
- log(`Thumbnail DB cleanup: folder=${folder} kept=${keptGalleryHashes.size} deleted=${delcount} time=${Math.round(t1 - t0)}ms`);
+ log('maintenanceQueue', { folder, kept: keptGalleryHashes.size, deleted: delcount, time: Math.round(t1 - t0) });
timer(`thumbnailDBCleanup:${folder}`, t1 - t0);
currentGalleryFolder = null;
el.clearCacheFolder.innerText = '';
updateStatusWithSort('Thumbnail cache cleared');
})
.catch((reason) => {
- SimpleFunctionQueue.abortLogger('Thumbnail DB cleanup:', reason);
+ SimpleFunctionQueue.abortLogger('thumbCacheCleanup', reason);
})
.finally(async () => {
await new Promise((resolve) => { setTimeout(resolve, 1000); }); // Delay removal by 1 second to ensure at least minimum visibility
@@ -1057,7 +1090,7 @@ function resetGalleryState(reason) {
function clearCacheIfDisabled(browser_cache) {
if (browser_cache === false) {
- log('Thumbnail DB cleanup:', 'Image gallery cache setting disabled. Clearing cache.');
+ log('thumbCacheCleanup', { disabled: true });
const controller = resetGalleryState('Clearing all thumbnails from cache');
maintenanceQueue.enqueue({
signal: controller.signal,
@@ -1066,13 +1099,13 @@ function clearCacheIfDisabled(browser_cache) {
const cb_clearMsg = showCleaningMsg(0, true);
await idbClearAll(controller.signal)
.then(() => {
- log(`Thumbnail DB cleanup: Cache cleared. time=${Math.floor(performance.now() - t0)}ms`);
+ log('thumbCacheCleanup', { time: Math.floor(performance.now() - t0) });
currentGalleryFolder = null;
el.clearCacheFolder.innerText = '';
updateStatusWithSort('Thumbnail cache cleared');
})
.catch((e) => {
- SimpleFunctionQueue.abortLogger('Thumbnail DB cleanup:', e);
+ SimpleFunctionQueue.abortLogger('thumbCacheCleanup', e);
})
.finally(async () => {
await new Promise((resolve) => { setTimeout(resolve, 1000); });
@@ -1307,6 +1340,80 @@ async function initGalleryAutoRefresh() {
galleryVisObserver.observe(galleryTab, { attributeFilter: ['class', 'style'], attributeOldValue: true });
}
+async function overlayDelete(evt) {
+ const res = await authFetch(`${window.api}/delete-image?file=${encodeURIComponent(currentImage)}`);
+ evt.stopPropagation();
+ if (!res || res.status !== 200) {
+ error('galleryDelete', { file: currentImage, status: res?.status, statusText: res?.statusText });
+ return;
+ }
+ const data = await res.json();
+ log('galleryDelete', data);
+ GalleryFolder.getActive()?.click();
+}
+
+async function overlayDownload(evt) {
+ log('galleryDownload', currentImage);
+ const link = document.createElement('a');
+ link.href = `/file=${encodeURIComponent(currentImage)}`;
+ link.download = currentName;
+ document.body.appendChild(link);
+ link.click();
+ document.body.removeChild(link);
+ evt.stopPropagation();
+}
+
+async function overlayInfo(evt) {
+ evt.stopPropagation();
+ const tgt = document.getElementById('html_info_formatted_gallery');
+ if (!tgt) return;
+ const res = await authFetch(`${window.api}/png-info?file=${encodeURI(currentImage)}`);
+ if (!res || res.status !== 200) return;
+ const data = await res.json();
+ log('galleryInfo res', data);
+ const prompt = data?.parameters?.Prompt || '';
+ const negative = data?.parameters?.Negative || data?.parameters?.['Negative prompt'] || '';
+ const raw = data?.info || '';
+ const params = data?.parameters || {};
+ delete params.Prompt;
+ delete params.Negative;
+ delete params['Negative prompt'];
+ const paramsFormatted = Object.entries(params).map(([key, value]) => `${key}: ${value}`).join(' | ');
+ tgt.innerHTML = `
+ File: ${currentImage}
+ Prompt: ${prompt}
+ Negative: ${negative}
+ ${paramsFormatted}
+
+ `;
+ const img = document.querySelector('#gallery_gallery img');
+ if (img) img.src = `/file=${encodeURIComponent(currentImage)}?t=${Date.now()}`; // Force refresh in case info endpoint is faster than cache update
+ const status = document.querySelector('#html_log_gallery p');
+ if (status) status.innerText = currentTitle;
+}
+
+async function createOverlay() {
+ if (el.overlay) return;
+ el.overlay = document.createElement('div');
+ el.overlay.className = 'gallery-overlay';
+ const btnDownload = document.createElement('span');
+ btnDownload.innerHTML = '\udb85\udc64';
+ btnDownload.title = 'Download image';
+ btnDownload.style.cursor = 'pointer';
+ btnDownload.addEventListener('click', overlayDownload);
+ const btnDelete = document.createElement('span');
+ btnDelete.innerHTML = '\uf05c';
+ btnDelete.title = 'Delete image';
+ btnDelete.style.cursor = 'pointer';
+ btnDelete.addEventListener('click', overlayDelete);
+ const btnInfo = document.createElement('span');
+ btnInfo.innerHTML = '\uf05a';
+ btnInfo.title = 'Image metadata';
+ btnInfo.style.cursor = 'pointer';
+ btnInfo.addEventListener('click', overlayInfo);
+ el.overlay.append(btnInfo, btnDelete, btnDownload);
+}
+
async function blockQueueUntilReady() {
// Add block to maintenanceQueue until cache is ready
maintenanceQueue.enqueue({
@@ -1329,22 +1436,27 @@ async function initGallery() { // triggered on gradio change to monitor when ui
el.files = gradioApp().getElementById('tab-gallery-files');
el.status = gradioApp().getElementById('tab-gallery-status');
el.search = gradioApp().querySelector('#tab-gallery-search textarea');
+ el.size = document.getElementById('tab-gallery-thumb-size');
if (!el.folders || !el.files || !el.status || !el.search) {
error('initGallery', 'Missing gallery elements');
return;
}
+ if (el.size) {
+ el.size.value = opts.extra_networks_card_size;
+ el.size.addEventListener('input', updateGalleryStyles);
+ }
blockQueueUntilReady(); // Run first
+ createOverlay();
updateGalleryStyles();
injectGalleryStatusCSS();
setOverlayAnimation();
galleryClearInit();
+
const progress = gradioApp().getElementById('tab-gallery-progress');
- if (progress) {
- galleryProgressBar.attachTo(progress);
- } else {
- log('initGallery', 'Failed to attach loading progress bar');
- }
+ if (progress) galleryProgressBar.attachTo(progress);
+ else log('initGallery', 'Failed to attach loading progress bar');
+
el.search.addEventListener('input', gallerySearch);
el.btnSend = gradioApp().getElementById('tab-gallery-send-image');
document.getElementById('tab-gallery-files').style.height = opts.logmonitor_show ? '75vh' : '85vh';
diff --git a/javascript/jquery.js b/javascript/jquery.js
new file mode 100644
index 000000000..c443b2fc9
--- /dev/null
+++ b/javascript/jquery.js
@@ -0,0 +1,2 @@
+/*! jQuery v3.7.0 -ajax,-ajax/jsonp,-ajax/load,-ajax/script,-ajax/var/location,-ajax/var/nonce,-ajax/var/rquery,-ajax/xhr,-manipulation/_evalUrl,-deprecated/ajax-event-alias,-effects,-effects/animatedSelector,-effects/Tween | (c) OpenJS Foundation and other contributors | jquery.org/license */
+!function(e,t){"use strict";"object"==typeof module&&"object"==typeof module.exports?module.exports=e.document?t(e,!0):function(e){if(!e.document)throw new Error("jQuery requires a window with a document");return t(e)}:t(e)}("undefined"!=typeof window?window:this,function(ie,e){"use strict";var oe=[],r=Object.getPrototypeOf,ae=oe.slice,g=oe.flat?function(e){return oe.flat.call(e)}:function(e){return oe.concat.apply([],e)},s=oe.push,se=oe.indexOf,n={},i=n.toString,ue=n.hasOwnProperty,o=ue.toString,a=o.call(Object),le={},v=function(e){return"function"==typeof e&&"number"!=typeof e.nodeType&&"function"!=typeof e.item},y=function(e){return null!=e&&e===e.window},m=ie.document,u={type:!0,src:!0,nonce:!0,noModule:!0};function b(e,t,n){var r,i,o=(n=n||m).createElement("script");if(o.text=e,t)for(r in u)(i=t[r]||t.getAttribute&&t.getAttribute(r))&&o.setAttribute(r,i);n.head.appendChild(o).parentNode.removeChild(o)}function x(e){return null==e?e+"":"object"==typeof e||"function"==typeof e?n[i.call(e)]||"object":typeof e}var t="3.7.0 -ajax,-ajax/jsonp,-ajax/load,-ajax/script,-ajax/var/location,-ajax/var/nonce,-ajax/var/rquery,-ajax/xhr,-manipulation/_evalUrl,-deprecated/ajax-event-alias,-effects,-effects/animatedSelector,-effects/Tween",l=/HTML$/i,ce=function(e,t){return new ce.fn.init(e,t)};function c(e){var t=!!e&&"length"in e&&e.length,n=x(e);return!v(e)&&!y(e)&&("array"===n||0===t||"number"==typeof t&&0+~]|"+ge+")"+ge+"*"),b=new RegExp(ge+"|>"),A=new RegExp(g),D=new RegExp("^"+t+"$"),N={ID:new RegExp("^#("+t+")"),CLASS:new RegExp("^\\.("+t+")"),TAG:new RegExp("^("+t+"|[*])"),ATTR:new RegExp("^"+d),PSEUDO:new RegExp("^"+g),CHILD:new RegExp("^:(only|first|last|nth|nth-last)-(child|of-type)(?:\\("+ge+"*(even|odd|(([+-]|)(\\d*)n|)"+ge+"*(?:([+-]|)"+ge+"*(\\d+)|))"+ge+"*\\)|)","i"),bool:new RegExp("^(?:"+f+")$","i"),needsContext:new RegExp("^"+ge+"*[>+~]|:(even|odd|eq|gt|lt|nth|first|last)(?:\\("+ge+"*((?:-\\d)?\\d*)"+ge+"*\\)|)(?=[^-]|$)","i")},L=/^(?:input|select|textarea|button)$/i,j=/^h\d$/i,O=/^(?:#([\w-]+)|(\w+)|\.([\w-]+))$/,P=/[+~]/,H=new RegExp("\\\\[\\da-fA-F]{1,6}"+ge+"?|\\\\([^\\r\\n\\f])","g"),q=function(e,t){var n="0x"+e.slice(1)-65536;return t||(n<0?String.fromCharCode(n+65536):String.fromCharCode(n>>10|55296,1023&n|56320))},R=function(){V()},M=K(function(e){return!0===e.disabled&&fe(e,"fieldset")},{dir:"parentNode",next:"legend"});try{E.apply(oe=ae.call(ye.childNodes),ye.childNodes),oe[ye.childNodes.length].nodeType}catch(e){E={apply:function(e,t){me.apply(e,ae.call(t))},call:function(e){me.apply(e,ae.call(arguments,1))}}}function I(t,e,n,r){var i,o,a,s,u,l,c,f=e&&e.ownerDocument,d=e?e.nodeType:9;if(n=n||[],"string"!=typeof t||!t||1!==d&&9!==d&&11!==d)return n;if(!r&&(V(e),e=e||T,C)){if(11!==d&&(u=O.exec(t)))if(i=u[1]){if(9===d){if(!(a=e.getElementById(i)))return n;if(a.id===i)return E.call(n,a),n}else if(f&&(a=f.getElementById(i))&&I.contains(e,a)&&a.id===i)return E.call(n,a),n}else{if(u[2])return E.apply(n,e.getElementsByTagName(t)),n;if((i=u[3])&&e.getElementsByClassName)return E.apply(n,e.getElementsByClassName(i)),n}if(!(h[t+" "]||p&&p.test(t))){if(c=t,f=e,1===d&&(b.test(t)||m.test(t))){(f=P.test(t)&&X(e.parentNode)||e)==e&&le.scope||((s=e.getAttribute("id"))?s=ce.escapeSelector(s):e.setAttribute("id",s=k)),o=(l=Y(t)).length;while(o--)l[o]=(s?"#"+s:":scope")+" "+G(l[o]);c=l.join(",")}try{return E.apply(n,f.querySelectorAll(c)),n}catch(e){h(t,!0)}finally{s===k&&e.removeAttribute("id")}}}return re(t.replace(ve,"$1"),e,n,r)}function W(){var r=[];return function e(t,n){return r.push(t+" ")>x.cacheLength&&delete e[r.shift()],e[t+" "]=n}}function B(e){return e[k]=!0,e}function F(e){var t=T.createElement("fieldset");try{return!!e(t)}catch(e){return!1}finally{t.parentNode&&t.parentNode.removeChild(t),t=null}}function $(t){return function(e){return fe(e,"input")&&e.type===t}}function _(t){return function(e){return(fe(e,"input")||fe(e,"button"))&&e.type===t}}function z(t){return function(e){return"form"in e?e.parentNode&&!1===e.disabled?"label"in e?"label"in e.parentNode?e.parentNode.disabled===t:e.disabled===t:e.isDisabled===t||e.isDisabled!==!t&&M(e)===t:e.disabled===t:"label"in e&&e.disabled===t}}function U(a){return B(function(o){return o=+o,B(function(e,t){var n,r=a([],e.length,o),i=r.length;while(i--)e[n=r[i]]&&(e[n]=!(t[n]=e[n]))})})}function X(e){return e&&"undefined"!=typeof e.getElementsByTagName&&e}function V(e){var t,n=e?e.ownerDocument||e:ye;return n!=T&&9===n.nodeType&&n.documentElement&&(r=(T=n).documentElement,C=!ce.isXMLDoc(T),i=r.matches||r.webkitMatchesSelector||r.msMatchesSelector,ye!=T&&(t=T.defaultView)&&t.top!==t&&t.addEventListener("unload",R),le.getById=F(function(e){return r.appendChild(e).id=ce.expando,!T.getElementsByName||!T.getElementsByName(ce.expando).length}),le.disconnectedMatch=F(function(e){return i.call(e,"*")}),le.scope=F(function(){return T.querySelectorAll(":scope")}),le.cssHas=F(function(){try{return T.querySelector(":has(*,:jqfake)"),!1}catch(e){return!0}}),le.getById?(x.filter.ID=function(e){var t=e.replace(H,q);return function(e){return e.getAttribute("id")===t}},x.find.ID=function(e,t){if("undefined"!=typeof t.getElementById&&C){var n=t.getElementById(e);return n?[n]:[]}}):(x.filter.ID=function(e){var n=e.replace(H,q);return function(e){var t="undefined"!=typeof e.getAttributeNode&&e.getAttributeNode("id");return t&&t.value===n}},x.find.ID=function(e,t){if("undefined"!=typeof t.getElementById&&C){var n,r,i,o=t.getElementById(e);if(o){if((n=o.getAttributeNode("id"))&&n.value===e)return[o];i=t.getElementsByName(e),r=0;while(o=i[r++])if((n=o.getAttributeNode("id"))&&n.value===e)return[o]}return[]}}),x.find.TAG=function(e,t){return"undefined"!=typeof t.getElementsByTagName?t.getElementsByTagName(e):t.querySelectorAll(e)},x.find.CLASS=function(e,t){if("undefined"!=typeof t.getElementsByClassName&&C)return t.getElementsByClassName(e)},p=[],F(function(e){var t;r.appendChild(e).innerHTML=" ",e.querySelectorAll("[selected]").length||p.push("\\["+ge+"*(?:value|"+f+")"),e.querySelectorAll("[id~="+k+"-]").length||p.push("~="),e.querySelectorAll("a#"+k+"+*").length||p.push(".#.+[+~]"),e.querySelectorAll(":checked").length||p.push(":checked"),(t=T.createElement("input")).setAttribute("type","hidden"),e.appendChild(t).setAttribute("name","D"),r.appendChild(e).disabled=!0,2!==e.querySelectorAll(":disabled").length&&p.push(":enabled",":disabled"),(t=T.createElement("input")).setAttribute("name",""),e.appendChild(t),e.querySelectorAll("[name='']").length||p.push("\\["+ge+"*name"+ge+"*="+ge+"*(?:''|\"\")")}),le.cssHas||p.push(":has"),p=p.length&&new RegExp(p.join("|")),l=function(e,t){if(e===t)return a=!0,0;var n=!e.compareDocumentPosition-!t.compareDocumentPosition;return n||(1&(n=(e.ownerDocument||e)==(t.ownerDocument||t)?e.compareDocumentPosition(t):1)||!le.sortDetached&&t.compareDocumentPosition(e)===n?e===T||e.ownerDocument==ye&&I.contains(ye,e)?-1:t===T||t.ownerDocument==ye&&I.contains(ye,t)?1:o?se.call(o,e)-se.call(o,t):0:4&n?-1:1)}),T}for(e in I.matches=function(e,t){return I(e,null,null,t)},I.matchesSelector=function(e,t){if(V(e),C&&!h[t+" "]&&(!p||!p.test(t)))try{var n=i.call(e,t);if(n||le.disconnectedMatch||e.document&&11!==e.document.nodeType)return n}catch(e){h(t,!0)}return 0":{dir:"parentNode",first:!0}," ":{dir:"parentNode"},"+":{dir:"previousSibling",first:!0},"~":{dir:"previousSibling"}},preFilter:{ATTR:function(e){return e[1]=e[1].replace(H,q),e[3]=(e[3]||e[4]||e[5]||"").replace(H,q),"~="===e[2]&&(e[3]=" "+e[3]+" "),e.slice(0,4)},CHILD:function(e){return e[1]=e[1].toLowerCase(),"nth"===e[1].slice(0,3)?(e[3]||I.error(e[0]),e[4]=+(e[4]?e[5]+(e[6]||1):2*("even"===e[3]||"odd"===e[3])),e[5]=+(e[7]+e[8]||"odd"===e[3])):e[3]&&I.error(e[0]),e},PSEUDO:function(e){var t,n=!e[6]&&e[2];return N.CHILD.test(e[0])?null:(e[3]?e[2]=e[4]||e[5]||"":n&&A.test(n)&&(t=Y(n,!0))&&(t=n.indexOf(")",n.length-t)-n.length)&&(e[0]=e[0].slice(0,t),e[2]=n.slice(0,t)),e.slice(0,3))}},filter:{TAG:function(e){var t=e.replace(H,q).toLowerCase();return"*"===e?function(){return!0}:function(e){return fe(e,t)}},CLASS:function(e){var t=s[e+" "];return t||(t=new RegExp("(^|"+ge+")"+e+"("+ge+"|$)"))&&s(e,function(e){return t.test("string"==typeof e.className&&e.className||"undefined"!=typeof e.getAttribute&&e.getAttribute("class")||"")})},ATTR:function(n,r,i){return function(e){var t=I.attr(e,n);return null==t?"!="===r:!r||(t+="","="===r?t===i:"!="===r?t!==i:"^="===r?i&&0===t.indexOf(i):"*="===r?i&&-1:\x20\t\r\n\f]*)[\x20\t\r\n\f]*\/?>(?:<\/\1>|)$/i;function C(e,n,r){return v(n)?ce.grep(e,function(e,t){return!!n.call(e,t,e)!==r}):n.nodeType?ce.grep(e,function(e){return e===n!==r}):"string"!=typeof n?ce.grep(e,function(e){return-1)[^>]*|#([\w-]+))$/;(ce.fn.init=function(e,t,n){var r,i;if(!e)return this;if(n=n||E,"string"==typeof e){if(!(r="<"===e[0]&&">"===e[e.length-1]&&3<=e.length?[null,e,null]:k.exec(e))||!r[1]&&t)return!t||t.jquery?(t||n).find(e):this.constructor(t).find(e);if(r[1]){if(t=t instanceof ce?t[0]:t,ce.merge(this,ce.parseHTML(r[1],t&&t.nodeType?t.ownerDocument||t:m,!0)),T.test(r[1])&&ce.isPlainObject(t))for(r in t)v(this[r])?this[r](t[r]):this.attr(r,t[r]);return this}return(i=m.getElementById(r[2]))&&(this[0]=i,this.length=1),this}return e.nodeType?(this[0]=e,this.length=1,this):v(e)?void 0!==n.ready?n.ready(e):e(ce):ce.makeArray(e,this)}).prototype=ce.fn,E=ce(m);var S=/^(?:parents|prev(?:Until|All))/,A={children:!0,contents:!0,next:!0,prev:!0};function D(e,t){while((e=e[t])&&1!==e.nodeType);return e}ce.fn.extend({has:function(e){var t=ce(e,this),n=t.length;return this.filter(function(){for(var e=0;e\x20\t\r\n\f]*)/i,Te=/^$|^module$|\/(?:java|ecma)script/i;re=m.createDocumentFragment().appendChild(m.createElement("div")),(be=m.createElement("input")).setAttribute("type","radio"),be.setAttribute("checked","checked"),be.setAttribute("name","t"),re.appendChild(be),le.checkClone=re.cloneNode(!0).cloneNode(!0).lastChild.checked,re.innerHTML="",le.noCloneChecked=!!re.cloneNode(!0).lastChild.defaultValue,re.innerHTML=" ",le.option=!!re.lastChild;var Ce={thead:[1,""],col:[2,""],tr:[2,""],td:[3,""],_default:[0,"",""]};function Ee(e,t){var n;return n="undefined"!=typeof e.getElementsByTagName?e.getElementsByTagName(t||"*"):"undefined"!=typeof e.querySelectorAll?e.querySelectorAll(t||"*"):[],void 0===t||t&&fe(e,t)?ce.merge([e],n):n}function ke(e,t){for(var n=0,r=e.length;n"," "]);var Se=/<|?\w+;/;function Ae(e,t,n,r,i){for(var o,a,s,u,l,c,f=t.createDocumentFragment(),d=[],p=0,h=e.length;p\s*$/g;function Re(e,t){return fe(e,"table")&&fe(11!==t.nodeType?t:t.firstChild,"tr")&&ce(e).children("tbody")[0]||e}function Me(e){return e.type=(null!==e.getAttribute("type"))+"/"+e.type,e}function Ie(e){return"true/"===(e.type||"").slice(0,5)?e.type=e.type.slice(5):e.removeAttribute("type"),e}function We(e,t){var n,r,i,o,a,s;if(1===t.nodeType){if(_.hasData(e)&&(s=_.get(e).events))for(i in _.remove(t,"handle events"),s)for(n=0,r=s[i].length;n",2===yt.childNodes.length),ce.parseHTML=function(e,t,n){return"string"!=typeof e?[]:("boolean"==typeof t&&(n=t,t=!1),t||(le.createHTMLDocument?((r=(t=m.implementation.createHTMLDocument("")).createElement("base")).href=m.location.href,t.head.appendChild(r)):t=m),o=!n&&[],(i=T.exec(e))?[t.createElement(i[1])]:(i=Ae([e],t,o),o&&o.length&&ce(o).remove(),ce.merge([],i.childNodes)));var r,i,o},ce.offset={setOffset:function(e,t,n){var r,i,o,a,s,u,l=ce.css(e,"position"),c=ce(e),f={};"static"===l&&(e.style.position="relative"),s=c.offset(),o=ce.css(e,"top"),u=ce.css(e,"left"),("absolute"===l||"fixed"===l)&&-1<(o+u).indexOf("auto")?(a=(r=c.position()).top,i=r.left):(a=parseFloat(o)||0,i=parseFloat(u)||0),v(t)&&(t=t.call(e,n,ce.extend({},s))),null!=t.top&&(f.top=t.top-s.top+a),null!=t.left&&(f.left=t.left-s.left+i),"using"in t?t.using.call(e,f):c.css(f)}},ce.fn.extend({offset:function(t){if(arguments.length)return void 0===t?this:this.each(function(e){ce.offset.setOffset(this,t,e)});var e,n,r=this[0];return r?r.getClientRects().length?(e=r.getBoundingClientRect(),n=r.ownerDocument.defaultView,{top:e.top+n.pageYOffset,left:e.left+n.pageXOffset}):{top:0,left:0}:void 0},position:function(){if(this[0]){var e,t,n,r=this[0],i={top:0,left:0};if("fixed"===ce.css(r,"position"))t=r.getBoundingClientRect();else{t=this.offset(),n=r.ownerDocument,e=r.offsetParent||n.documentElement;while(e&&(e===n.body||e===n.documentElement)&&"static"===ce.css(e,"position"))e=e.parentNode;e&&e!==r&&1===e.nodeType&&((i=ce(e).offset()).top+=ce.css(e,"borderTopWidth",!0),i.left+=ce.css(e,"borderLeftWidth",!0))}return{top:t.top-i.top-ce.css(r,"marginTop",!0),left:t.left-i.left-ce.css(r,"marginLeft",!0)}}},offsetParent:function(){return this.map(function(){var e=this.offsetParent;while(e&&"static"===ce.css(e,"position"))e=e.offsetParent;return e||K})}}),ce.each({scrollLeft:"pageXOffset",scrollTop:"pageYOffset"},function(t,i){var o="pageYOffset"===i;ce.fn[t]=function(e){return R(this,function(e,t,n){var r;if(y(e)?r=e:9===e.nodeType&&(r=e.defaultView),void 0===n)return r?r[i]:e[t];r?r.scrollTo(o?r.pageXOffset:n,o?n:r.pageYOffset):e[t]=n},t,e,arguments.length)}}),ce.each(["top","left"],function(e,n){ce.cssHooks[n]=Qe(le.pixelPosition,function(e,t){if(t)return t=Ve(e,n),$e.test(t)?ce(e).position()[n]+"px":t})}),ce.each({Height:"height",Width:"width"},function(a,s){ce.each({padding:"inner"+a,content:s,"":"outer"+a},function(r,o){ce.fn[o]=function(e,t){var n=arguments.length&&(r||"boolean"!=typeof e),i=r||(!0===e||!0===t?"margin":"border");return R(this,function(e,t,n){var r;return y(e)?0===o.indexOf("outer")?e["inner"+a]:e.document.documentElement["client"+a]:9===e.nodeType?(r=e.documentElement,Math.max(e.body["scroll"+a],r["scroll"+a],e.body["offset"+a],r["offset"+a],r["client"+a])):void 0===n?ce.css(e,t,i):ce.style(e,t,n,i)},s,n?e:void 0,n)}})}),ce.fn.extend({bind:function(e,t,n){return this.on(e,null,t,n)},unbind:function(e,t){return this.off(e,null,t)},delegate:function(e,t,n,r){return this.on(t,e,n,r)},undelegate:function(e,t,n){return 1===arguments.length?this.off(e,"**"):this.off(t,e||"**",n)},hover:function(e,t){return this.mouseenter(e).mouseleave(t||e)}}),ce.each("blur focus focusin focusout resize scroll click dblclick mousedown mouseup mousemove mouseover mouseout mouseenter mouseleave change select submit keydown keypress keyup contextmenu".split(" "),function(e,n){ce.fn[n]=function(e,t){return 0 {
+ const removeLivePreview = (ok = false) => {
debug('taskEnd:', id_task);
localStorage.removeItem('task');
setProgress();
@@ -124,6 +126,11 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
for (const gallery of galleries) gallery.style.display = 'flex'; // remove all galleries
try {
if (parentGallery && livePreview) {
+ if (ok) {
+ const previewImg = gradioApp().querySelector('#livePreviewImage');
+ const galleryImg = gradioApp().querySelector('#control_gallery img');
+ if (previewImg?.src && galleryImg) galleryImg.src = previewImg.src; // copy preview to gallery if everything is ok
+ }
parentGallery.removeChild(livePreview);
parentGallery.style.minHeight = 'unset';
parentGallery.style.maxHeight = 'unset';
@@ -135,18 +142,28 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
if (atEnd) atEnd();
};
- const start = (id_task, id_live_preview) => { // eslint-disable-line no-shadow
+ const startLivePreview = (id_task, id_live_preview) => { // eslint-disable-line no-shadow
if (opts.live_preview_refresh_period === 0) return;
const request_id = document.hidden ? -1 : id_live_preview;
const onProgressHandler = (res) => {
- if (res?.debug) debug('livePreview:', dateStart, request_id, res);
+ if (res?.debug) debug('progress:', { start: dateStart, id: request_id, res });
lastState = res;
const elapsedFromStart = (new Date() - dateStart) / 1000;
hasStarted |= res.active;
- if (res.completed || (!res.active && (hasStarted || once)) || (elapsedFromStart > 120 && !res.queued && res.progress === prevProgress)) {
- debug('livePreview end:', res);
- done();
+ if (res.completed || (!res.active && (hasStarted || once))) {
+ debug('progress', { end: res, reason: res.completed ? 'completed' : 'inactive' });
+ if (!res.paused) removeLivePreview(true); // only abort if not paused
+ return;
+ }
+ if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) {
+ debug('progress', { end: res, reason: 'progressSimeout' });
+ if (!res.paused) removeLivePreview(false); // only abort if not paused
+ return;
+ }
+ if (elapsedFromStart > startTimeout && !res.queued && !res.active) {
+ debug('progress', { end: res, reason: 'startTimeout' });
+ if (!res.paused) removeLivePreview(false); // only abort if not paused
return;
}
if (res.progress !== prevProgress) {
@@ -160,16 +177,16 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
id_live_preview = res.id_live_preview;
}
if (onProgress) onProgress(res);
- setTimeout(() => start(id_task, id_live_preview), opts.live_preview_refresh_period || 500);
+ setTimeout(() => startLivePreview(id_task, id_live_preview), opts.live_preview_refresh_period || 500);
};
const onProgressErrorHandler = (err) => {
- error(`livePreview: ${err}`);
- done();
+ error('progress', { error: err });
+ removeLivePreview(false);
};
xhrPost('./internal/progress', { id_task, id_live_preview: request_id }, onProgressHandler, onProgressErrorHandler, false, 30000);
};
- debug('livePreview start:', dateStart);
- start(id_task, 0);
+ debug('progress', { start: dateStart });
+ startLivePreview(id_task, 0);
}
diff --git a/javascript/script.js b/javascript/script.js
index c4428584e..f462642e1 100644
--- a/javascript/script.js
+++ b/javascript/script.js
@@ -41,12 +41,17 @@ const optionsChangedCallbacks = [];
let uiCurrentTab = null;
let uiAfterUpdateTimeout = null;
+function registerCallback(queue, callback) {
+ if (queue.includes(callback)) return;
+ queue.push(callback);
+}
+
function onAfterUiUpdate(callback) {
if (typeof callback !== 'function') {
error(`onAfterUiUpdate was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- uiAfterUpdateCallbacks.push(callback);
+ registerCallback(uiAfterUpdateCallbacks, callback);
}
function onUiUpdate(callback) {
@@ -54,7 +59,7 @@ function onUiUpdate(callback) {
error(`onUiUpdate was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- uiUpdateCallbacks.push(callback);
+ registerCallback(uiUpdateCallbacks, callback);
}
function onUiLoaded(callback) {
@@ -62,7 +67,7 @@ function onUiLoaded(callback) {
error(`onUiLoaded was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- uiLoadedCallbacks.push(callback);
+ registerCallback(uiLoadedCallbacks, callback);
}
function onUiReady(callback) {
@@ -70,7 +75,7 @@ function onUiReady(callback) {
error(`onUiReady was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- uiReadyCallbacks.push(callback);
+ registerCallback(uiReadyCallbacks, callback);
}
function onUiTabChange(callback) {
@@ -78,7 +83,7 @@ function onUiTabChange(callback) {
error(`onUiTabChange was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- uiTabChangeCallbacks.push(callback);
+ registerCallback(uiTabChangeCallbacks, callback);
}
function onOptionsChanged(callback) {
@@ -86,7 +91,7 @@ function onOptionsChanged(callback) {
error(`onOptionsChanged was called without a valid value. Expected a function but got: ${callback}`);
return;
}
- optionsChangedCallbacks.push(callback);
+ registerCallback(optionsChangedCallbacks, callback);
}
function executeCallbacks(queue, arg) {
@@ -179,6 +184,88 @@ document.addEventListener('keydown', (e) => {
}
});
+function getSortableCellValue(cell, sortType) {
+ const rawValue = cell?.dataset?.sortValue ?? cell?.textContent?.trim() ?? '';
+ if (sortType === 'number') {
+ const numericValue = Number.parseFloat(rawValue);
+ return Number.isNaN(numericValue) ? Number.NEGATIVE_INFINITY : numericValue;
+ }
+ return rawValue.toLowerCase();
+}
+
+function sortTable(table, columnIndex, sortType, sortOrder) {
+ const tbody = table.querySelector('tbody');
+ if (!tbody) return;
+ const rows = Array.from(tbody.querySelectorAll('tr'));
+ const direction = sortOrder === 'desc' ? -1 : 1;
+ const sortedRows = rows
+ .map((row, index) => ({ row, index }))
+ .sort((a, b) => {
+ const aCell = a.row.children[columnIndex];
+ const bCell = b.row.children[columnIndex];
+ const aValue = getSortableCellValue(aCell, sortType);
+ const bValue = getSortableCellValue(bCell, sortType);
+ if (aValue < bValue) return -1 * direction;
+ if (aValue > bValue) return 1 * direction;
+ return a.index - b.index;
+ });
+ tbody.replaceChildren(...sortedRows.map((item) => item.row));
+}
+
+function applySortIndicators(table, activeHeader, sortOrder) {
+ const headers = table.querySelectorAll('th.sortable');
+ for (const header of headers) {
+ header.classList.remove('sorted-asc', 'sorted-desc');
+ header.removeAttribute('aria-sort');
+ }
+ activeHeader.classList.add(sortOrder === 'desc' ? 'sorted-desc' : 'sorted-asc');
+ activeHeader.setAttribute('aria-sort', sortOrder === 'desc' ? 'descending' : 'ascending');
+}
+
+function handleSortableTableClick(event) {
+ const header = event.target.closest('th.sortable');
+ if (!header) return;
+ const table = header.closest('table[data-sortable="true"]');
+ if (!table) return;
+ const headers = Array.from(table.querySelectorAll('th.sortable'));
+ const columnIndex = headers.indexOf(header);
+ if (columnIndex < 0) return;
+
+ const currentSortKey = table.dataset.sortKey || table.dataset.defaultSortKey;
+ const currentSortOrder = table.dataset.sortOrder || table.dataset.defaultSortOrder || 'asc';
+ const isCurrentHeader = currentSortKey === header.dataset.sortKey;
+ const nextOrder = isCurrentHeader && currentSortOrder === 'asc' ? 'desc' : 'asc';
+
+ table.dataset.sortKey = header.dataset.sortKey;
+ table.dataset.sortOrder = nextOrder;
+ sortTable(table, columnIndex, header.dataset.sortType || 'text', nextOrder);
+ applySortIndicators(table, header, nextOrder);
+}
+
+async function initTableSorter() {
+ const t0 = performance.now();
+ const root = gradioApp();
+ if (!root.dataset.tableSorterBound) {
+ root.addEventListener('click', handleSortableTableClick);
+ root.dataset.tableSorterBound = 'true';
+ }
+ const t1 = performance.now();
+ log('initTableSorter', Math.round(t1 - t0));
+ timer('initTableSorter', t1 - t0);
+}
+
+async function deleteFile(filename) {
+ if (!filename) return;
+ if (!confirm(`Are you sure you want to delete the object - This action cannot be undone? Object: ${filename}`)) return; // eslint-disable-line no-alert
+ const res = await authFetch(`${window.api}/delete-file?file=${encodeURIComponent(filename)}`);
+ if (!res || res.status !== 200) {
+ error('FileDelete', { file: filename, status: res?.status, statusText: res?.statusText });
+ return;
+ }
+ const data = await res.json();
+ log('FileDelete', data);
+}
+
/**
* checks that a UI element is not in another hidden element or tab content
*/
diff --git a/javascript/sdnext.css b/javascript/sdnext.css
index e06ed40e2..1bf73a8cd 100644
--- a/javascript/sdnext.css
+++ b/javascript/sdnext.css
@@ -2193,6 +2193,30 @@ div:has(>#tab-gallery-folders) {
background-color: var(--button-primary-border-color) !important;
}
+.simple-table th.sortable {
+ cursor: pointer;
+ user-select: none;
+ position: relative;
+ padding-right: 1.2em;
+}
+
+.simple-table th.sortable::after {
+ content: '↕';
+ position: absolute;
+ right: 0.3em;
+ opacity: 0.55;
+}
+
+.simple-table th.sortable.sorted-asc::after {
+ content: '↑';
+ opacity: 1;
+}
+
+.simple-table th.sortable.sorted-desc::after {
+ content: '↓';
+ opacity: 1;
+}
+
.simple-table tr:nth-child(odd) {
background-color: var(--neutral-900);
}
diff --git a/javascript/sparkline.js b/javascript/sparkline.js
new file mode 100644
index 000000000..e227771f0
--- /dev/null
+++ b/javascript/sparkline.js
@@ -0,0 +1,5 @@
+/* jquery.sparkline 2.1.2 - http://omnipotent.net/jquery.sparkline/
+** Licensed under the New BSD License - see above site for details */
+
+(function(a,b,c){(function(a){typeof define=="function"&&define.amd?define(["jquery"],a):jQuery&&!jQuery.fn.sparkline&&a(jQuery)})(function(d){"use strict";var e={},f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,A,B,C,D,E,F,G,H,I,J,K,L=0;f=function(){return{common:{type:"line",lineColor:"#00f",fillColor:"#cdf",defaultPixelsPerValue:3,width:"auto",height:"auto",composite:!1,tagValuesAttribute:"values",tagOptionsPrefix:"spark",enableTagOptions:!1,enableHighlight:!0,highlightLighten:1.4,tooltipSkipNull:!0,tooltipPrefix:"",tooltipSuffix:"",disableHiddenCheck:!1,numberFormatter:!1,numberDigitGroupCount:3,numberDigitGroupSep:",",numberDecimalMark:".",disableTooltips:!1,disableInteraction:!1},line:{spotColor:"#f80",highlightSpotColor:"#5f5",highlightLineColor:"#f22",spotRadius:1.5,minSpotColor:"#f80",maxSpotColor:"#f80",lineWidth:1,normalRangeMin:c,normalRangeMax:c,normalRangeColor:"#ccc",drawNormalOnTop:!1,chartRangeMin:c,chartRangeMax:c,chartRangeMinX:c,chartRangeMaxX:c,tooltipFormat:new h('● {{prefix}}{{y}}{{suffix}}')},bar:{barColor:"#3366cc",negBarColor:"#f44",stackedBarColor:["#3366cc","#dc3912","#ff9900","#109618","#66aa00","#dd4477","#0099c6","#990099"],zeroColor:c,nullColor:c,zeroAxis:!0,barWidth:4,barSpacing:1,chartRangeMax:c,chartRangeMin:c,chartRangeClip:!1,colorMap:c,tooltipFormat:new h('● {{prefix}}{{value}}{{suffix}}')},tristate:{barWidth:4,barSpacing:1,posBarColor:"#6f6",negBarColor:"#f44",zeroBarColor:"#999",colorMap:{},tooltipFormat:new h('● {{value:map}}'),tooltipValueLookups:{map:{"-1":"Loss",0:"Draw",1:"Win"}}},discrete:{lineHeight:"auto",thresholdColor:c,thresholdValue:0,chartRangeMax:c,chartRangeMin:c,chartRangeClip:!1,tooltipFormat:new h("{{prefix}}{{value}}{{suffix}}")},bullet:{targetColor:"#f33",targetWidth:3,performanceColor:"#33f",rangeColors:["#d3dafe","#a8b6ff","#7f94ff"],base:c,tooltipFormat:new h("{{fieldkey:fields}} - {{value}}"),tooltipValueLookups:{fields:{r:"Range",p:"Performance",t:"Target"}}},pie:{offset:0,sliceColors:["#3366cc","#dc3912","#ff9900","#109618","#66aa00","#dd4477","#0099c6","#990099"],borderWidth:0,borderColor:"#000",tooltipFormat:new h('● {{value}} ({{percent.1}}%)')},box:{raw:!1,boxLineColor:"#000",boxFillColor:"#cdf",whiskerColor:"#000",outlierLineColor:"#333",outlierFillColor:"#fff",medianColor:"#f00",showOutliers:!0,outlierIQR:1.5,spotRadius:1.5,target:c,targetColor:"#4a2",chartRangeMax:c,chartRangeMin:c,tooltipFormat:new h("{{field:fields}}: {{value}}"),tooltipFormatFieldlistKey:"field",tooltipValueLookups:{fields:{lq:"Lower Quartile",med:"Median",uq:"Upper Quartile",lo:"Left Outlier",ro:"Right Outlier",lw:"Left Whisker",rw:"Right Whisker"}}}}},E='.jqstooltip { position: absolute;left: 0px;top: 0px;visibility: hidden;background: rgb(0, 0, 0) transparent;background-color: rgba(0,0,0,0.6);filter:progid:DXImageTransform.Microsoft.gradient(startColorstr=#99000000, endColorstr=#99000000);-ms-filter: "progid:DXImageTransform.Microsoft.gradient(startColorstr=#99000000, endColorstr=#99000000)";color: white;font: 10px arial, san serif;text-align: left;white-space: nowrap;padding: 5px;border: 1px solid white;z-index: 10000;}.jqsfield { color: white;font: 10px arial, san serif;text-align: left;}',g=function(){var a,b;return a=function(){this.init.apply(this,arguments)},arguments.length>1?(arguments[0]?(a.prototype=d.extend(new arguments[0],arguments[arguments.length-1]),a._super=arguments[0].prototype):a.prototype=arguments[arguments.length-1],arguments.length>2&&(b=Array.prototype.slice.call(arguments,1,-1),b.unshift(a.prototype),d.extend.apply(d,b))):a.prototype=arguments[0],a.prototype.cls=a,a},d.SPFormatClass=h=g({fre:/\{\{([\w.]+?)(:(.+?))?\}\}/g,precre:/(\w+)\.(\d+)/,init:function(a,b){this.format=a,this.fclass=b},render:function(a,b,d){var e=this,f=a,g,h,i,j,k;return this.format.replace(this.fre,function(){var a;return h=arguments[1],i=arguments[3],g=e.precre.exec(h),g?(k=g[2],h=g[1]):k=!1,j=f[h],j===c?"":i&&b&&b[i]?(a=b[i],a.get?b[i].get(j)||j:b[i][j]||j):(n(j)&&(d.get("numberFormatter")?j=d.get("numberFormatter")(j):j=s(j,k,d.get("numberDigitGroupCount"),d.get("numberDigitGroupSep"),d.get("numberDecimalMark"))),j)})}}),d.spformat=function(a,b){return new h(a,b)},i=function(a,b,c){return ac?c:a},j=function(a,c){var d;return c===2?(d=b.floor(a.length/2),a.length%2?a[d]:(a[d-1]+a[d])/2):a.length%2?(d=(a.length*c+c)/4,d%1?(a[b.floor(d)]+a[b.floor(d)-1])/2:a[d-1]):(d=(a.length*c+2)/4,d%1?(a[b.floor(d)]+a[b.floor(d)-1])/2:a[d-1])},k=function(a){var b;switch(a){case"undefined":a=c;break;case"null":a=null;break;case"true":a=!0;break;case"false":a=!1;break;default:b=parseFloat(a),a==b&&(a=b)}return a},l=function(a){var b,c=[];for(b=a.length;b--;)c[b]=k(a[b]);return c},m=function(a,b){var c,d,e=[];for(c=0,d=a.length;c0;h-=c)a.splice(h,0,e);return a.join("")},o=function(a,b,c){var d;for(d=b.length;d--;){if(c&&b[d]===null)continue;if(b[d]!==a)return!1}return!0},p=function(a){var b=0,c;for(c=a.length;c--;)b+=typeof a[c]=="number"?a[c]:0;return b},r=function(a){return d.isArray(a)?a:[a]},q=function(b){var c;a.createStyleSheet?a.createStyleSheet().cssText=b:(c=a.createElement("style"),c.type="text/css",a.getElementsByTagName("head")[0].appendChild(c),c[typeof a.body.style.WebkitAppearance=="string"?"innerText":"innerHTML"]=b)},d.fn.simpledraw=function(b,e,f,g){var h,i;if(f&&(h=this.data("_jqs_vcanvas")))return h;if(d.fn.sparkline.canvas===!1)return!1;if(d.fn.sparkline.canvas===c){var j=a.createElement("canvas");if(!j.getContext||!j.getContext("2d")){if(!a.namespaces||!!a.namespaces.v)return d.fn.sparkline.canvas=!1,!1;a.namespaces.add("v","urn:schemas-microsoft-com:vml","#default#VML"),d.fn.sparkline.canvas=function(a,b,c,d){return new J(a,b,c)}}else d.fn.sparkline.canvas=function(a,b,c,d){return new I(a,b,c,d)}}return b===c&&(b=d(this).innerWidth()),e===c&&(e=d(this).innerHeight()),h=d.fn.sparkline.canvas(b,e,this,g),i=d(this).data("_jqs_mhandler"),i&&i.registerCanvas(h),h},d.fn.cleardraw=function(){var a=this.data("_jqs_vcanvas");a&&a.reset()},d.RangeMapClass=t=g({init:function(a){var b,c,d=[];for(b in a)a.hasOwnProperty(b)&&typeof b=="string"&&b.indexOf(":")>-1&&(c=b.split(":"),c[0]=c[0].length===0?-Infinity:parseFloat(c[0]),c[1]=c[1].length===0?Infinity:parseFloat(c[1]),c[2]=a[b],d.push(c));this.map=a,this.rangelist=d||!1},get:function(a){var b=this.rangelist,d,e,f;if((f=this.map[a])!==c)return f;if(b)for(d=b.length;d--;){e=b[d];if(e[0]<=a&&e[1]>=a)return e[2]}return c}}),d.range_map=function(a){return new t(a)},u=g({init:function(a,b){var c=d(a);this.$el=c,this.options=b,this.currentPageX=0,this.currentPageY=0,this.el=a,this.splist=[],this.tooltip=null,this.over=!1,this.displayTooltips=!b.get("disableTooltips"),this.highlightEnabled=!b.get("disableHighlight")},registerSparkline:function(a){this.splist.push(a),this.over&&this.updateDisplay()},registerCanvas:function(a){var b=d(a.canvas);this.canvas=a,this.$canvas=b,b.mouseenter(d.proxy(this.mouseenter,this)),b.mouseleave(d.proxy(this.mouseleave,this)),b.click(d.proxy(this.mouseclick,this))},reset:function(a){this.splist=[],this.tooltip&&a&&(this.tooltip.remove(),this.tooltip=c)},mouseclick:function(a){var b=d.Event("sparklineClick");b.originalEvent=a,b.sparklines=this.splist,this.$el.trigger(b)},mouseenter:function(b){d(a.body).unbind("mousemove.jqs"),d(a.body).bind("mousemove.jqs",d.proxy(this.mousemove,this)),this.over=!0,this.currentPageX=b.pageX,this.currentPageY=b.pageY,this.currentEl=b.target,!this.tooltip&&this.displayTooltips&&(this.tooltip=new v(this.options),this.tooltip.updatePosition(b.pageX,b.pageY)),this.updateDisplay()},mouseleave:function(){d(a.body).unbind("mousemove.jqs");var b=this.splist,c=b.length,e=!1,f,g;this.over=!1,this.currentEl=null,this.tooltip&&(this.tooltip.remove(),this.tooltip=null);for(g=0;g",{id:"jqssizetip",style:e,"class":c}),this.tooltip=d("
",{id:"jqstooltip","class":c}).appendTo(this.container),f=this.tooltip.offset(),this.offsetLeft=f.left,this.offsetTop=f.top,this.hidden=!0,d(window).unbind("resize.jqs scroll.jqs"),d(window).bind("resize.jqs scroll.jqs",d.proxy(this.updateWindowDims,this)),this.updateWindowDims()},updateWindowDims:function(){this.scrollTop=d(window).scrollTop(),this.scrollLeft=d(window).scrollLeft(),this.scrollRight=this.scrollLeft+d(window).width(),this.updatePosition()},getSize:function(a){this.sizetip.html(a).appendTo(this.container),this.width=this.sizetip.width()+1,this.height=this.sizetip.height(),this.sizetip.remove()},setContent:function(a){if(!a){this.tooltip.css("visibility","hidden"),this.hidden=!0;return}this.getSize(a),this.tooltip.html(a).css({width:this.width,height:this.height,visibility:"visible"}),this.hidden&&(this.hidden=!1,this.updatePosition())},updatePosition:function(a,b){if(a===c){if(this.mousex===c)return;a=this.mousex-this.offsetLeft,b=this.mousey-this.offsetTop}else this.mousex=a-=this.offsetLeft,this.mousey=b-=this.offsetTop;if(!this.height||!this.width||this.hidden)return;b-=this.height+this.tooltipOffsetY,a+=this.tooltipOffsetX,bthis.scrollRight&&(a=this.scrollRight-this.width),this.tooltip.css({left:a,top:b})},remove:function(){this.tooltip.remove(),this.sizetip.remove(),this.sizetip=this.tooltip=c,d(window).unbind("resize.jqs scroll.jqs")}}),F=function(){q(E)},d(F),K=[],d.fn.sparkline=function(b,e){return this.each(function(){var f=new d.fn.sparkline.options(this,e),g=d(this),h,i;h=function(){var e,h,i,j,k,l,m;if(b==="html"||b===c){m=this.getAttribute(f.get("tagValuesAttribute"));if(m===c||m===null)m=g.html();e=m.replace(/(^\s*\s*$)|\s+/g,"").split(",")}else e=b;h=f.get("width")==="auto"?e.length*f.get("defaultPixelsPerValue"):f.get("width");if(f.get("height")==="auto"){if(!f.get("composite")||!d.data(this,"_jqs_vcanvas"))j=a.createElement("span"),j.innerHTML="a",g.html(j),i=d(j).innerHeight()||d(j).height(),d(j).remove(),j=null}else i=f.get("height");f.get("disableInteraction")?k=!1:(k=d.data(this,"_jqs_mhandler"),k?f.get("composite")||k.reset():(k=new u(this,f),d.data(this,"_jqs_mhandler",k)));if(f.get("composite")&&!d.data(this,"_jqs_vcanvas")){d.data(this,"_jqs_errnotify")||(alert("Attempted to attach a composite sparkline to an element with no existing sparkline"),d.data(this,"_jqs_errnotify",!0));return}l=new(d.fn.sparkline[f.get("type")])(this,e,f,h,i),l.render(),k&&k.registerSparkline(l)};if(d(this).html()&&!f.get("disableHiddenCheck")&&d(this).is(":hidden")||!d(this).parents("body").length){if(!f.get("composite")&&d.data(this,"_jqs_pending"))for(i=K.length;i;i--)K[i-1][0]==this&&K.splice(i-1,1);K.push([this,h]),d.data(this,"_jqs_pending",!0)}else h.call(this)})},d.fn.sparkline.defaults=f(),d.sparkline_display_visible=function(){var a,b,c,e=[];for(b=0,c=K.length;bthis.canvasWidth||d>this.canvasHeight||b<0||d<0?null:(g=this.getRegion(a,b,d),e!==g?(e!==c&&f&&this.removeHighlight(),this.currentRegion=g,g!==c&&f&&this.renderHighlight(),!0):!1)},clearRegionHighlight:function(){return this.currentRegion!==c?(this.removeHighlight(),this.currentRegion=c,!0):!1},renderHighlight:function(){this.changeHighlight(!0)},removeHighlight:function(){this.changeHighlight(!1)},changeHighlight:function(a){},getCurrentRegionTooltip:function(){var a=this.options,b="",e=[],f,g,i,j,k,l,m,n,o,p,q,r,s,t;if(this.currentRegion===c)return"";f=this.getCurrentRegionFields(),q=a.get("tooltipFormatter");if(q)return q(this,a,f);a.get("tooltipChartTitle")&&(b+=''+a.get("tooltipChartTitle")+"
\n"),g=this.options.get("tooltipFormat");if(!g)return"";d.isArray(g)||(g=[g]),d.isArray(f)||(f=[f]),m=this.options.get("tooltipFormatFieldlist"),n=this.options.get("tooltipFormatFieldlistKey");if(m&&n){o=[];for(l=f.length;l--;)p=f[l][n],(t=d.inArray(p,m))!=-1&&(o[t]=f[l]);f=o}i=g.length,s=f.length;for(l=0;l'+k+"")}return e.length?b+e.join("\n"):""},getCurrentRegionFields:function(){},calcHighlightColor:function(a,c){var d=c.get("highlightColor"),e=c.get("highlightLighten"),f,g,h,j;if(d)return d;if(e){f=/^#([0-9a-f])([0-9a-f])([0-9a-f])$/i.exec(a)||/^#([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})$/i.exec(a);if(f){h=[],g=a.length===4?16:1;for(j=0;j<3;j++)h[j]=i(b.round(parseInt(f[j+1],16)*g*e),0,255);return"rgb("+h.join(",")+")"}}return a}}),w={changeHighlight:function(a){var b=this.currentRegion,c=this.target,e=this.regionShapes[b],f;e&&(f=this.renderRegion(b,a),d.isArray(f)||d.isArray(e)?(c.replaceWithShapes(e,f),this.regionShapes[b]=d.map(f,function(a){return a.id})):(c.replaceWithShape(e,f),this.regionShapes[b]=f.id))},render:function(){var a=this.values,b=this.target,c=this.regionShapes,e,f,g,h;if(!this.cls._super.render.call(this))return;for(g=a.length;g--;){e=this.renderRegion(g);if(e)if(d.isArray(e)){f=[];for(h=e.length;h--;)e[h].append(),f.push(e[h].id);c[g]=f}else e.append(),c[g]=e.id;else c[g]=null}b.render()}},d.fn.sparkline.line=x=g(d.fn.sparkline._base,{type:"line",init:function(a,b,c,d,e){x._super.init.call(this,a,b,c,d,e),this.vertices=[],this.regionMap=[],this.xvalues=[],this.yvalues=[],this.yminmax=[],this.hightlightSpotId=null,this.lastShapeId=null,this.initTarget()},getRegion:function(a,b,d){var e,f=this.regionMap;for(e=f.length;e--;)if(f[e]!==null&&b>=f[e][0]&&b<=f[e][1])return f[e][2];return c},getCurrentRegionFields:function(){var a=this.currentRegion;return{isNull:this.yvalues[a]===null,x:this.xvalues[a],y:this.yvalues[a],color:this.options.get("lineColor"),fillColor:this.options.get("fillColor"),offset:a}},renderHighlight:function(){var a=this.currentRegion,b=this.target,d=this.vertices[a],e=this.options,f=e.get("spotRadius"),g=e.get("highlightSpotColor"),h=e.get("highlightLineColor"),i,j;if(!d)return;f&&g&&(i=b.drawCircle(d[0],d[1],f,c,g),this.highlightSpotId=i.id,b.insertAfterShape(this.lastShapeId,i)),h&&(j=b.drawLine(d[0],this.canvasTop,d[0],this.canvasTop+this.canvasHeight,h),this.highlightLineId=j.id,b.insertAfterShape(this.lastShapeId,j))},removeHighlight:function(){var a=this.target;this.highlightSpotId&&(a.removeShapeId(this.highlightSpotId),this.highlightSpotId=null),this.highlightLineId&&(a.removeShapeId(this.highlightLineId),this.highlightLineId=null)},scanValues:function(){var a=this.values,c=a.length,d=this.xvalues,e=this.yvalues,f=this.yminmax,g,h,i,j,k;for(g=0;gthis.maxy&&(this.maxy=d)),a.get("chartRangeMin")!==c&&(a.get("chartRangeClip")||a.get("chartRangeMin")this.maxy)&&(this.maxy=a.get("chartRangeMax")),a.get("chartRangeMinX")!==c&&(a.get("chartRangeClipX")||a.get("chartRangeMinX")this.maxx)&&(this.maxx=a.get("chartRangeMaxX"))},drawNormalRange:function(a,d,e,f,g){var h=this.options.get("normalRangeMin"),i=this.options.get("normalRangeMax"),j=d+b.round(e-e*((i-this.miny)/g)),k=b.round(e*(i-h)/g);this.target.drawRect(a,j,f,k,c,this.options.get("normalRangeColor")).append()},render:function(){var a=this.options,e=this.target,f=this.canvasWidth,g=this.canvasHeight,h=this.vertices,i=a.get("spotRadius"),j=this.regionMap,k,l,m,n,o,p,q,r,s,u,v,w,y,z,A,B,C,D,E,F,G,H,I,J,K;if(!x._super.render.call(this))return;this.scanValues(),this.processRangeOptions(),I=this.xvalues,J=this.yvalues;if(!this.yminmax.length||this.yvalues.length<2)return;n=o=0,k=this.maxx-this.minx===0?1:this.maxx-this.minx,l=this.maxy-this.miny===0?1:this.maxy-this.miny,m=this.yvalues.length-1,i&&(fthis.maxy&&(u=this.maxy),q.length||q.push([w,n+g]),p=[w,n+b.round(g-g*((u-this.miny)/l))],q.push(p),h.push(p));C=[],D=[],E=r.length;for(K=0;K2&&(q[0]=[q[0][0],q[1][1]]),C.push(q));E=D.length;for(K=0;K-1;if(u||d.isArray(O))J=!0,u&&(O=e[A]=l(O.split(":"))),O=m(O,null),v=b.min.apply(b,O),w=b.max.apply(b,O),vs&&(s=w)}this.stacked=J,this.regionShapes={},this.barWidth=j,this.barSpacing=n,this.totalBarWidth=j+n,this.width=g=e.length*j+(e.length-1)*n,this.initTarget(),q&&(H=o===c?-Infinity:o,I=p===c?Infinity:p),z=[],x=J?[]:z;var S=[],T=[];for(A=0,B=e.length;A0&&(S[A]+=O),r<0&&s>0?O<0?T[A]+=b.abs(O):x[A]+=O:x[A]+=b.abs(O-(O<0?s:r)),z.push(O))}else O=q?i(e[A],H,I):e[A],O=e[A]=k(O),O!==null&&z.push(O);this.max=G=b.max.apply(b,z),this.min=F=b.min.apply(b,z),this.stackMax=s=J?b.max.apply(b,S):G,this.stackMin=r=J?b.min.apply(b,z):F,f.get("chartRangeMin")!==c&&(f.get("chartRangeClip")||f.get("chartRangeMin")G)&&(G=f.get("chartRangeMax")),this.zeroAxis=D=f.get("zeroAxis",!0),F<=0&&G>=0&&D?E=0:D==0?E=F:F>0?E=F:E=G,this.xaxisOffset=E,C=J?b.max.apply(b,x)+b.max.apply(b,T):G-F,this.canvasHeightEf=D&&F<0?this.canvasHeight-2:this.canvasHeight-1,F=0?s:G,P=(Q-E)/C*this.canvasHeight,P!==b.ceil(P)&&(this.canvasHeightEf-=2,P=b.ceil(P))):P=this.canvasHeight,this.yoffset=P,d.isArray(f.get("colorMap"))?(this.colorMapByIndex=f.get("colorMap"),this.colorMapByValue=null):(this.colorMapByIndex=null,this.colorMapByValue=f.get("colorMap"),this.colorMapByValue&&this.colorMapByValue.get===c&&(this.colorMapByValue=new t(this.colorMapByValue))),this.range=C},getRegion:function(a,d,e){var f=b.floor(d/this.totalBarWidth);return f<0||f>=this.values.length?c:f},getCurrentRegionFields:function(){var a=this.currentRegion,b=r(this.values[a]),c=[],d,e;for(e=b.length;e--;)d=b[e],c.push({isNull:d===null,value:d,color:this.calcColor(e,d,a),offset:a});return c},calcColor:function(a,b,e){var f=this.colorMapByIndex,g=this.colorMapByValue,h=this.options,i,j;return this.stacked?i=h.get("stackedBarColor"):i=b<0?h.get("negBarColor"):h.get("barColor"),b===0&&h.get("zeroColor")!==c&&(i=h.get("zeroColor")),g&&(j=g.get(b))?i=j:f&&f.length>e&&(i=f[e]),d.isArray(i)?i[a%i.length]:i},renderRegion:function(a,e){var f=this.values[a],g=this.options,h=this.xaxisOffset,i=[],j=this.range,k=this.stacked,l=this.target,m=a*this.totalBarWidth,n=this.canvasHeightEf,p=this.yoffset,q,r,s,t,u,v,w,x,y,z;f=d.isArray(f)?f:[f],w=f.length,x=f[0],t=o(null,f),z=o(h,f,!0);if(t)return g.get("nullColor")?(s=e?g.get("nullColor"):this.calcHighlightColor(g.get("nullColor"),g),q=p>0?p-1:p,l.drawRect(m,q,this.barWidth-1,0,s,s)):c;u=p;for(v=0;v0?r=b.floor(n*(b.abs(x-h)/j))+1:r=1,xb?g=e[b]:c[b]<0?g=d.get("negBarColor"):c[b]>0?g=d.get("posBarColor"):g=d.get("zeroBarColor"),g},renderRegion:function(a,c){var d=this.values,e=this.options,f=this.target,g,h,i,j,k,l;g=f.pixelHeight,i=b.round(g/2),j=a*this.totalBarWidth,d[a]<0?(k=i,h=i-1):d[a]>0?(k=0,h=i-1):(k=i-1,h=2),l=this.calcColor(d[a],a);if(l===null)return;return c&&(l=this.calcHighlightColor(l,e)),f.drawRect(j,k,this.barWidth-1,h-1,l,l)}}),d.fn.sparkline.discrete=A=g(d.fn.sparkline._base,w,{type:"discrete",init:function(a,e,f,g,h){A._super.init.call(this,a,e,f,g,h),this.regionShapes={},this.values=e=d.map(e,Number),this.min=b.min.apply(b,e),this.max=b.max.apply(b,e),this.range=this.max-this.min,this.width=g=f.get("width")==="auto"?e.length*2:this.width,this.interval=b.floor(g/e.length),this.itemWidth=g/e.length,f.get("chartRangeMin")!==c&&(f.get("chartRangeClip")||f.get("chartRangeMin")this.max)&&(this.max=f.get("chartRangeMax")),this.initTarget(),this.target&&(this.lineHeight=f.get("lineHeight")==="auto"?b.round(this.canvasHeight*.3):f.get("lineHeight"))},getRegion:function(a,c,d){return b.floor(c/this.itemWidth)},getCurrentRegionFields:function(){var a=this.currentRegion;return{isNull:this.values[a]===c,value:this.values[a],offset:a}},renderRegion:function(a,c){var d=this.values,e=this.options,f=this.min,g=this.max,h=this.range,j=this.interval,k=this.target,l=this.canvasHeight,m=this.lineHeight,n=l-m,o,p,q,r;return p=i(d[a],f,g),r=a*j,o=b.round(n-n*((p-f)/h)),q=e.get("thresholdColor")&&p0)for(i=c.length;i--;)h+=c[i];this.total=h,this.initTarget(),this.radius=b.floor(b.min(this.canvasWidth,this.canvasHeight)/2)},getRegion:function(a,b,d){var e=this.target.getShapeAt(a,b,d);return e!==c&&this.shapes[e]!==c?this.shapes[e]:c},getCurrentRegionFields:function(){var a=this.currentRegion;return{isNull:this.values[a]===c,value:this.values[a],percent:this.values[a]/this.total*100,color:this.options.get("sliceColors")[a%this.options.get("sliceColors").length],offset:a}},changeHighlight:function(a){var b=this.currentRegion,c=this.renderSlice(b,a),d=this.valueShapes[b];delete this.shapes[d],this.target.replaceWithShape(d,c),this.valueShapes[b]=c.id,this.shapes[c.id]=b},renderSlice:function(a,d){var e=this.target,f=this.options,g=this.radius,h=f.get("borderWidth"),i=f.get("offset"),j=2*b.PI,k=this.values,l=this.total,m=i?2*b.PI*(i/360):0,n,o,p,q,r;q=k.length;for(p=0;p0&&(o=m+j*(k[p]/l));if(a===p)return r=f.get("sliceColors")[p%f.get("sliceColors").length],d&&(r=this.calcHighlightColor(r,f)),e.drawPieSlice(g,g,g-h,n,o,c,r);m=o}},render:function(){var a=this.target,d=this.values,e=this.options,f=this.radius,g=e.get("borderWidth"),h,i;if(!C._super.render.call(this))return;g&&a.drawCircle(f,f,b.floor(f-g/2),e.get("borderColor"),c,g).append();for(i=d.length;i--;)d[i]&&(h=this.renderSlice(i).append(),this.valueShapes[i]=h.id,this.shapes[h.id]=i);a.render()}}),d.fn.sparkline.box=D=g(d.fn.sparkline._base,{type:"box",init:function(a,b,c,e,f){D._super.init.call(this,a,b,c,e,f),this.values=d.map(b,Number),this.width=c.get("width")==="auto"?"4.0em":e,this.initTarget(),this.values.length||(this.disabled=1)},getRegion:function(){return 1},getCurrentRegionFields:function(){var a=[{field:"lq",value:this.quartiles[0]},{field:"med",value:this.quartiles
+[1]},{field:"uq",value:this.quartiles[2]}];return this.loutlier!==c&&a.push({field:"lo",value:this.loutlier}),this.routlier!==c&&a.push({field:"ro",value:this.routlier}),this.lwhisker!==c&&a.push({field:"lw",value:this.lwhisker}),this.rwhisker!==c&&a.push({field:"rw",value:this.rwhisker}),a},render:function(){var a=this.target,d=this.values,e=d.length,f=this.options,g=this.canvasWidth,h=this.canvasHeight,i=f.get("chartRangeMin")===c?b.min.apply(b,d):f.get("chartRangeMin"),k=f.get("chartRangeMax")===c?b.max.apply(b,d):f.get("chartRangeMax"),l=0,m,n,o,p,q,r,s,t,u,v,w;if(!D._super.render.call(this))return;if(f.get("raw"))f.get("showOutliers")&&d.length>5?(n=d[0],m=d[1],p=d[2],q=d[3],r=d[4],s=d[5],t=d[6]):(m=d[0],p=d[1],q=d[2],r=d[3],s=d[4]);else{d.sort(function(a,b){return a-b}),p=j(d,1),q=j(d,2),r=j(d,3),o=r-p;if(f.get("showOutliers")){m=s=c;for(u=0;up-o*f.get("outlierIQR")&&(m=d[u]),d[u]s&&a.drawCircle((t-i)*w+l,h/2,f.get("spotRadius"),f.get("outlierLineColor"),f.get("outlierFillColor")).append()),a.drawRect(b.round((p-i)*w+l),b.round(h*.1),b.round((r-p)*w),b.round(h*.8),f.get("boxLineColor"),f.get("boxFillColor")).append(),a.drawLine(b.round((m-i)*w+l),b.round(h/2),b.round((p-i)*w+l),b.round(h/2),f.get("lineColor")).append(),a.drawLine(b.round((m-i)*w+l),b.round(h/4),b.round((m-i)*w+l),b.round(h-h/4),f.get("whiskerColor")).append(),a.drawLine(b.round((s-i)*w+l),b.round(h/2),b.round((r-i)*w+l),b.round(h/2),f.get("lineColor")).append(),a.drawLine(b.round((s-i)*w+l),b.round(h/4),b.round((s-i)*w+l),b.round(h-h/4),f.get("whiskerColor")).append(),a.drawLine(b.round((q-i)*w+l),b.round(h*.1),b.round((q-i)*w+l),b.round(h*.9),f.get("medianColor")).append(),f.get("target")&&(v=b.ceil(f.get("spotRadius")),a.drawLine(b.round((f.get("target")-i)*w+l),b.round(h/2-v),b.round((f.get("target")-i)*w+l),b.round(h/2+v),f.get("targetColor")).append(),a.drawLine(b.round((f.get("target")-i)*w+l-v),b.round(h/2),b.round((f.get("target")-i)*w+l+v),b.round(h/2),f.get("targetColor")).append()),a.render()}}),G=g({init:function(a,b,c,d){this.target=a,this.id=b,this.type=c,this.args=d},append:function(){return this.target.appendShape(this),this}}),H=g({_pxregex:/(\d+)(px)?\s*$/i,init:function(a,b,c){if(!a)return;this.width=a,this.height=b,this.target=c,this.lastShapeId=null,c[0]&&(c=c[0]),d.data(c,"_jqs_vcanvas",this)},drawLine:function(a,b,c,d,e,f){return this.drawShape([[a,b],[c,d]],e,f)},drawShape:function(a,b,c,d){return this._genShape("Shape",[a,b,c,d])},drawCircle:function(a,b,c,d,e,f){return this._genShape("Circle",[a,b,c,d,e,f])},drawPieSlice:function(a,b,c,d,e,f,g){return this._genShape("PieSlice",[a,b,c,d,e,f,g])},drawRect:function(a,b,c,d,e,f){return this._genShape("Rect",[a,b,c,d,e,f])},getElement:function(){return this.canvas},getLastShapeId:function(){return this.lastShapeId},reset:function(){alert("reset not implemented")},_insert:function(a,b){d(b).html(a)},_calculatePixelDims:function(a,b,c){var e;e=this._pxregex.exec(b),e?this.pixelHeight=e[1]:this.pixelHeight=d(c).height(),e=this._pxregex.exec(a),e?this.pixelWidth=e[1]:this.pixelWidth=d(c).width()},_genShape:function(a,b){var c=L++;return b.unshift(c),new G(this,c,a,b)},appendShape:function(a){alert("appendShape not implemented")},replaceWithShape:function(a,b){alert("replaceWithShape not implemented")},insertAfterShape:function(a,b){alert("insertAfterShape not implemented")},removeShapeId:function(a){alert("removeShapeId not implemented")},getShapeAt:function(a,b,c){alert("getShapeAt not implemented")},render:function(){alert("render not implemented")}}),I=g(H,{init:function(b,e,f,g){I._super.init.call(this,b,e,f),this.canvas=a.createElement("canvas"),f[0]&&(f=f[0]),d.data(f,"_jqs_vcanvas",this),d(this.canvas).css({display:"inline-block",width:b,height:e,verticalAlign:"top"}),this._insert(this.canvas,f),this._calculatePixelDims(b,e,this.canvas),this.canvas.width=this.pixelWidth,this.canvas.height=this.pixelHeight,this.interact=g,this.shapes={},this.shapeseq=[],this.currentTargetShapeId=c,d(this.canvas).css({width:this.pixelWidth,height:this.pixelHeight})},_getContext:function(a,b,d){var e=this.canvas.getContext("2d");return a!==c&&(e.strokeStyle=a),e.lineWidth=d===c?1:d,b!==c&&(e.fillStyle=b),e},reset:function(){var a=this._getContext();a.clearRect(0,0,this.pixelWidth,this.pixelHeight),this.shapes={},this.shapeseq=[],this.currentTargetShapeId=c},_drawShape:function(a,b,d,e,f){var g=this._getContext(d,e,f),h,i;g.beginPath(),g.moveTo(b[0][0]+.5,b[0][1]+.5);for(h=1,i=b.length;h',this.canvas.insertAdjacentHTML("beforeEnd",f),this.group=d(this.canvas).children()[0],this.rendered=!1,this.prerender=""},_drawShape:function(a,b,d,e,f){var g=[],h,i,j,k,l,m,n;for(n=0,m=b.length;n'+" ",l},_drawCircle:function(a,b,d,e,f,g,h){var i,j,k;return b-=e,d-=e,i=f===c?' stroked="false" ':' strokeWeight="'+h+'px" strokeColor="'+f+'" ',j=g===c?' filled="false"':' fillColor="'+g+'" filled="true" ',k=' ',k},_drawPieSlice:function(a,d,e,f,g,h,i,j){var k,l,m,n,o,p,q,r;if(g===h)return"";h-g===2*b.PI&&(g=0,h=2*b.PI),l=d+b.round(b.cos(g)*f),m=e+b.round(b.sin(g)*f),n=d+b.round(b.cos(h)*f),o=e+b.round(b.sin(h)*f);if(l===n&&m===o){if(h-g'+" ",r)},_drawRect:function(a,b,c,d,e,f,g){return this._drawShape(a,[[b,c],[b,c+e],[b+d,c+e],[b+d,c],[b,c]],f,g)},reset:function(){this.group.innerHTML=""},appendShape:function(a){var b=this["_draw"+a.type].apply(this,a.args);return this.rendered?this.group.insertAdjacentHTML("beforeEnd",b):this.prerender+=b,this.lastShapeId=a.id,a.id},replaceWithShape:function(a,b){var c=d("#jqsshape"+a),e=this["_draw"+b.type].apply(this,b.args);c[0].outerHTML=e},replaceWithShapes:function(a,b){var c=d("#jqsshape"+a[0]),e="",f=b.length,g;for(g=0;g b[1] - a[1]);
- const filteredTimers = allTimers.filter((t) => t[1] > 50);
- debug('startupTimers', filteredTimers);
- // xhrPost(`${window.api}/log`, { debug: JSON.stringify(filteredTimers) });
+ // allTimers.sort((a, b) => b[1] - a[1]);
+ const filteredTimers = allTimers.filter((t) => t[1] > 100);
+ const objTimers = {};
+ for (const [name, elapsed] of filteredTimers) objTimers[name] = elapsed;
+ debug('startupTimers', objTimers);
+ // xhrPost(`${window.api}/log`, { debug: JSON.stringify(objTimers) });
}
window.timer = timer;
diff --git a/javascript/ui.js b/javascript/ui.js
index 0c973cc10..1cf6ee42f 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -108,10 +108,6 @@ function send_to_kanvas(gallery) {
const [image] = extract_image_from_gallery(gallery);
log('sendToKanvas', image);
if (window.loadFromURL && image.data) window.loadFromURL(image.data);
- // const inputPanelEl = gradioApp().getElementById('control-template-column-input');
- // if (inputPanelEl) inputPanelEl.classList.remove('hidden');
- const inputPanelCb = gradioApp().getElementById('control_dynamic_input');
- if (inputPanelCb && !inputPanelCb.checked) inputPanelCb.click();
}
async function setTheme(val, old) {
@@ -638,6 +634,7 @@ function selectCheckpoint(name) {
else gradioApp().getElementById('change_checkpoint').click();
log(`selectCheckpoint ${isRefiner ? 'refiner' : 'model'}: ${desiredCheckpointName}`);
markSelectedCards([desiredCheckpointName], 'model');
+ setTimeout(requestProgress, 250);
}
let desiredVAEName = null;
@@ -661,6 +658,7 @@ function selectReference(name) {
desiredCheckpointName = name;
gradioApp().getElementById('change_reference').click();
markSelectedCards([desiredCheckpointName], 'model');
+ setTimeout(requestProgress, 250);
}
function currentImageResolutionimg2img(_a, _b, scaleBy) {
diff --git a/models/Reference/HiDream-ai--HiDream-O1-Image-Dev.jpg b/models/Reference/HiDream-ai--HiDream-O1-Image-Dev.jpg
new file mode 100644
index 000000000..c7661f473
Binary files /dev/null and b/models/Reference/HiDream-ai--HiDream-O1-Image-Dev.jpg differ
diff --git a/models/Reference/HiDream-ai--HiDream-O1-Image.jpg b/models/Reference/HiDream-ai--HiDream-O1-Image.jpg
new file mode 100644
index 000000000..c8274c7b2
Binary files /dev/null and b/models/Reference/HiDream-ai--HiDream-O1-Image.jpg differ
diff --git a/models/Reference/Owen777--UltraFlux-v1.jpg b/models/Reference/Owen777--UltraFlux-v1.jpg
new file mode 100644
index 000000000..c81ec159f
Binary files /dev/null and b/models/Reference/Owen777--UltraFlux-v1.jpg differ
diff --git a/models/Reference/jdopensource--JoyAI-Image-Edit-Diffusers.jpg b/models/Reference/jdopensource--JoyAI-Image-Edit-Diffusers.jpg
new file mode 100644
index 000000000..ba1c7c984
Binary files /dev/null and b/models/Reference/jdopensource--JoyAI-Image-Edit-Diffusers.jpg differ
diff --git a/models/Reference/stepfun-ai--Step1X-Edit-v1p1-diffusers.jpg b/models/Reference/stepfun-ai--Step1X-Edit-v1p1-diffusers.jpg
new file mode 100644
index 000000000..c4ebba557
Binary files /dev/null and b/models/Reference/stepfun-ai--Step1X-Edit-v1p1-diffusers.jpg differ
diff --git a/models/Reference/vladmandic--VIBE-Image-Edit.jpg b/models/Reference/vladmandic--VIBE-Image-Edit.jpg
new file mode 100644
index 000000000..0eff72f2f
Binary files /dev/null and b/models/Reference/vladmandic--VIBE-Image-Edit.jpg differ
diff --git a/modules/api/api.py b/modules/api/api.py
index 098d83f98..0eb6c30e9 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -93,6 +93,10 @@ class Api:
self.add_api_route("/sdapi/v1/unets", endpoints.get_unets, methods=["GET"], response_model=list[models.ItemUNet])
# functional api
+ self.add_api_route("/sdapi/v1/file", endpoints.get_file, methods=["GET"], tags=["Functional"])
+ self.add_api_route("/sdapi/v1/delete-image", endpoints.get_deleteimage, methods=["GET"], tags=["Functional"])
+ self.add_api_route("/sdapi/v1/delete-file", endpoints.get_deletefile, methods=["GET"], tags=["Functional"])
+ self.add_api_route("/sdapi/v1/png-info", endpoints.get_pnginfo, methods=["GET"], response_model=models.ResImageInfo, tags=["Functional"])
self.add_api_route("/sdapi/v1/png-info", endpoints.post_pnginfo, methods=["POST"], response_model=models.ResImageInfo, tags=["Functional"])
self.add_api_route("/sdapi/v1/checkpoint", endpoints.get_checkpoint, methods=["GET"], tags=["Functional"])
self.add_api_route("/sdapi/v1/checkpoint", endpoints.set_checkpoint, methods=["POST"], tags=["Functional"])
diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py
index 2d27c6003..eec9f3666 100644
--- a/modules/api/endpoints.py
+++ b/modules/api/endpoints.py
@@ -1,4 +1,6 @@
+from fastapi.exceptions import HTTPException
from modules import shared
+from modules.logger import log
from modules.api import models, helpers
@@ -323,6 +325,95 @@ def get_extensions_list():
})
return ext_list
+def get_file(file: str):
+ import os
+ from pathlib import Path
+ from starlette.responses import FileResponse
+ allowed_dirs = shared.demo.allowed_paths
+ if not file.strip():
+ raise HTTPException(status_code=400, detail="file path is required")
+ if not any(Path(folder).absolute() in Path(file).absolute().parents for folder in allowed_dirs):
+ raise HTTPException(status_code=403, detail=f"file {file}: must be in one of allowed directories")
+ if not os.path.exists(file):
+ raise HTTPException(status_code=404, detail=f"file not found: {file}")
+ if os.path.isdir(file):
+ raise HTTPException(status_code=403, detail=f"file {file}: is a directory")
+ return FileResponse(file, media_type='application/octet-stream', filename=file)
+
+def get_deletefile(file: str):
+ import os
+ from pathlib import Path
+ allowed_dirs = shared.demo.allowed_paths
+ if file is None or len(file.strip()) == 0:
+ raise HTTPException(status_code=400, detail="file path is required")
+ if not any(Path(folder).absolute() in Path(file).absolute().parents for folder in allowed_dirs):
+ raise HTTPException(status_code=403, detail=f"file {file}: must be in one of allowed directories")
+ if not os.path.exists(file):
+ raise HTTPException(status_code=404, detail=f"file not found: {file}")
+ try:
+ if os.path.isdir(file):
+ log.warning(f'Delete: folder="{file}"')
+ import shutil
+ shutil.rmtree(file)
+ else:
+ log.warning(f'Delete: file="{file}"')
+ os.remove(file)
+ return {"deleted": f"{file}"}
+ except Exception as e:
+ log.error(f'Delete: file="{file}" error: {e}')
+ raise HTTPException(status_code=500, detail=f"error deleting file {file}: {str(e)}") from e
+
+def get_deleteimage(file: str):
+ import os
+ from pathlib import Path
+ allowed_dirs = shared.demo.allowed_paths
+ if file is None or len(file.strip()) == 0:
+ raise HTTPException(status_code=400, detail="file path is required")
+ if not any(Path(folder).absolute() in Path(file).absolute().parents for folder in allowed_dirs):
+ raise HTTPException(status_code=403, detail=f"file {file}: must be in one of allowed directories")
+ if not os.path.exists(file):
+ raise HTTPException(status_code=404, detail=f"file not found: {file}")
+ if os.path.isdir(file):
+ raise HTTPException(status_code=403, detail=f"file {file}: is a directory")
+ if os.path.splitext(file)[1].lower() not in (".png", ".jpg", ".jpeg", ".webp"):
+ raise HTTPException(status_code=403, detail=f"file {file}: not an image file")
+ try:
+ os.remove(file)
+ log.warning(f'Delete: image="{file}"')
+ return {"deleted": f"{file}"}
+ except Exception as e:
+ log.error(f'Delete: file="{file}" error: {e}')
+ raise HTTPException(status_code=500, detail=f"error deleting file {file}: {str(e)}") from e
+
+def get_pnginfo(file: str):
+ """Extract generation parameters from a image file path. Returns raw info string and parsed parameters dict."""
+ import os
+ from pathlib import Path
+ from PIL import Image
+ from modules import images, infotext
+ allowed_dirs = shared.demo.allowed_paths
+ if not file.strip():
+ raise HTTPException(status_code=400, detail="file path is required")
+ if not any(Path(folder).absolute() in Path(file).absolute().parents for folder in allowed_dirs):
+ raise HTTPException(status_code=403, detail=f"file {file}: must be in one of allowed directories")
+ if os.path.splitext(file)[1].lower() not in (".png", ".jpg", ".jpeg", ".webp"):
+ raise HTTPException(status_code=403, detail=f"file {file}: not an image file")
+ if not os.path.isfile(file):
+ raise HTTPException(status_code=403, detail=f"file {file}: not an image file")
+ image = None
+ try:
+ image = Image.open(file)
+ image.load()
+ except Exception as e:
+ raise HTTPException(status_code=403, detail=f"file {file}: not an image file") from e
+ if image is None:
+ raise HTTPException(status_code=403, detail=f"file {file}: not an image file")
+ geninfo, items = images.read_info_from_image(image)
+ if geninfo is None:
+ geninfo = ""
+ params = infotext.parse(geninfo)
+ return models.ResImageInfo(info=geninfo, items=items, parameters=params)
+
def post_pnginfo(req: models.ReqImageInfo):
"""Extract generation parameters from a PNG image's metadata. Returns raw info string and parsed parameters dict."""
from modules import images, script_callbacks, infotext
diff --git a/modules/api/server.py b/modules/api/server.py
index 7daf91af5..db1e75f8d 100644
--- a/modules/api/server.py
+++ b/modules/api/server.py
@@ -116,8 +116,8 @@ def get_progress(req: models.ReqProgress = Depends()):
progress = min((current / total) if current > 0 and total > 0 else 0, 1)
time_since_start = time.time() - shared.state.time_start
eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
- # log.critical(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
- # log.critical(shared.state)
+ # log.trace(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
+ # log.trace(shared.state)
res = models.ResProgress(id=shared.state.id, progress=round(progress, 2), eta_relative=round(eta_relative, 2), current_image=current_image, textinfo=shared.state.textinfo, state=shared.state.dict(), )
return res
diff --git a/modules/api/validate.py b/modules/api/validate.py
index 58952332c..1ff77be75 100644
--- a/modules/api/validate.py
+++ b/modules/api/validate.py
@@ -5,6 +5,7 @@ from modules.logger import log
# value is cost: -1=disabled, 0=unlimited, 1=default, >1 expensive
request_cost = {
"/file": 0,
+ "/internal/progress": 0,
"/run/predict": 0,
"/sdapi/v1/browser/thumb": 0,
"/sdapi/v1/network/thumb": 0,
diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py
index 5ec6aa5b3..710c3b7e3 100644
--- a/modules/caption/joytag.py
+++ b/modules/caption/joytag.py
@@ -1,6 +1,5 @@
# Vendored from JoyTag: https://huggingface.co/spaces/fancyfeast/joytag
# Contains full model architecture (ViT, CNN stems, MAE) including training-only code
-# retained for update compatibility. Do not modify directly — sync from upstream.
import os
import math
@@ -942,7 +941,7 @@ class ViT(VisionModel):
loss_type: str,
layerscale_init: float | None = None,
head_mean_after: bool = False,
- cnn_stem: str = None,
+ cnn_stem: str | None = None,
patch_dropout: float = 0.0,
):
super().__init__(image_size, n_tags)
diff --git a/modules/cfgzero/cogview4_pipeline.py b/modules/cfgzero/cogview4_pipeline.py
index 472c2eb72..07c6a8383 100644
--- a/modules/cfgzero/cogview4_pipeline.py
+++ b/modules/cfgzero/cogview4_pipeline.py
@@ -191,7 +191,7 @@ class CogView4CFGZeroPipeline(DiffusionPipeline, CogView4LoraLoaderMixin):
def _get_glm_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
max_sequence_length: int = 1024,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
diff --git a/modules/cfgzero/flux_pipeline.py b/modules/cfgzero/flux_pipeline.py
index 362a434e4..e42d7baa1 100644
--- a/modules/cfgzero/flux_pipeline.py
+++ b/modules/cfgzero/flux_pipeline.py
@@ -217,7 +217,7 @@ class FluxCFGZeroPipeline(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
device: Optional[torch.device] = None,
@@ -535,7 +535,7 @@ class FluxCFGZeroPipeline(
@staticmethod
def _unpack_latents(latents, height, width, vae_scale_factor):
- batch_size, num_patches, channels = latents.shape
+ batch_size, _num_patches, channels = latents.shape
# VAE applies 8x compression on images but we must also account for packing which requires
# latent height and width to be divisible by 2.
@@ -637,9 +637,9 @@ class FluxCFGZeroPipeline(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ negative_prompt: Union[str, List[str]] | None = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
true_cfg_scale: float = 1.0,
height: Optional[int] = None,
diff --git a/modules/cfgzero/hidream_pipeline.py b/modules/cfgzero/hidream_pipeline.py
index 9eaff0878..edfd47d78 100644
--- a/modules/cfgzero/hidream_pipeline.py
+++ b/modules/cfgzero/hidream_pipeline.py
@@ -211,7 +211,7 @@ class HiDreamImageCFGZeroPipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
max_sequence_length: int = 128,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
@@ -285,7 +285,7 @@ class HiDreamImageCFGZeroPipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin
def _get_llama3_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
max_sequence_length: int = 128,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
@@ -545,7 +545,7 @@ class HiDreamImageCFGZeroPipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
prompt_3: Optional[Union[str, List[str]]] = None,
prompt_4: Optional[Union[str, List[str]]] = None,
diff --git a/modules/cfgzero/hunyuan_t2v_pipeline.py b/modules/cfgzero/hunyuan_t2v_pipeline.py
index 8494b9b96..0f880f576 100644
--- a/modules/cfgzero/hunyuan_t2v_pipeline.py
+++ b/modules/cfgzero/hunyuan_t2v_pipeline.py
@@ -317,7 +317,7 @@ class HunyuanVideoCFGZeroPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin
def encode_prompt(
self,
prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]] = None,
+ prompt_2: Union[str, List[str]] | None = None,
prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
num_videos_per_prompt: int = 1,
prompt_embeds: Optional[torch.Tensor] = None,
@@ -481,15 +481,15 @@ class HunyuanVideoCFGZeroPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
+ prompt_2: Union[str, List[str]] | None = None,
+ negative_prompt: Union[str, List[str]] | None = None,
+ negative_prompt_2: Union[str, List[str]] | None = None,
height: int = 720,
width: int = 1280,
num_frames: int = 129,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: List[float] | None = None,
true_cfg_scale: float = 1.0,
guidance_scale: float = 6.0,
num_videos_per_prompt: Optional[int] = 1,
diff --git a/modules/cfgzero/sd3_pipeline.py b/modules/cfgzero/sd3_pipeline.py
index 946571d81..1cae06511 100644
--- a/modules/cfgzero/sd3_pipeline.py
+++ b/modules/cfgzero/sd3_pipeline.py
@@ -246,7 +246,7 @@ class StableDiffusion3CFGZeroPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
device: Optional[torch.device] = None,
@@ -786,7 +786,7 @@ class StableDiffusion3CFGZeroPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
prompt_3: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
@@ -813,7 +813,7 @@ class StableDiffusion3CFGZeroPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 256,
- skip_guidance_layers: List[int] = None,
+ skip_guidance_layers: List[int] | None = None,
skip_layer_guidance_scale: float = 2.8,
skip_layer_guidance_stop: float = 0.2,
skip_layer_guidance_start: float = 0.01,
diff --git a/modules/cfgzero/wan_t2v_pipeline.py b/modules/cfgzero/wan_t2v_pipeline.py
index 9cc0fa529..4660b2360 100644
--- a/modules/cfgzero/wan_t2v_pipeline.py
+++ b/modules/cfgzero/wan_t2v_pipeline.py
@@ -153,7 +153,7 @@ class WanCFGZeroPipeline(DiffusionPipeline, WanLoraLoaderMixin):
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
device: Optional[torch.device] = None,
@@ -374,8 +374,8 @@ class WanCFGZeroPipeline(DiffusionPipeline, WanLoraLoaderMixin):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
+ negative_prompt: Union[str, List[str]] | None = None,
height: int = 480,
width: int = 832,
num_frames: int = 81,
diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py
index 154b93448..4e3f05525 100644
--- a/modules/civitai/download_civitai.py
+++ b/modules/civitai/download_civitai.py
@@ -323,10 +323,13 @@ class DownloadManager:
if version and version.images:
for img in version.images:
if img.url:
- code, _size, _note = download_civit_preview(final_file, img.url)
+ code, _size, _note = download_civit_preview(final_file, img.url, meta=img.meta)
if code == 200:
log.info(f'CivitAI preview saved: id={item.id}')
break
+ if code == 304 and backfill_preview_parameters(final_file, img.url, img.meta):
+ log.info(f'CivitAI preview backfilled: id={item.id}')
+ break
except Exception as e:
log.warning(f'CivitAI preview fetch failed: id={item.id} {e}')
@@ -340,6 +343,306 @@ class DownloadManager:
download_manager = DownloadManager()
+# ---- Preview metadata helpers ----
+
+NOISE_META_KEYS = frozenset({'hashes', 'comfy', 'comfyui', 'workflow', 'extrametadata'})
+PASSTHROUGH_VALUE_LIMIT = 512 # chars; pass-through values longer than this are dropped
+
+
+def civitai_meta_to_parameters(meta: dict | None) -> str:
+ """Convert Civitai version-image meta dict to sdnext parameters string.
+
+ Output matches sdnext's standard `parameters` channel (`modules/image/save.py:65-72`):
+ positive prompt on line one, optional `Negative prompt:` line two,
+ comma-joined `Key: Value` pairs on line three. Round-trippable through
+ `modules.infotext.parse`. Generator-specific noise keys (full ComfyUI
+ workflows, etc.) are dropped and any pass-through value exceeding
+ `PASSTHROUGH_VALUE_LIMIT` chars is omitted to keep the embedded chunk
+ compact.
+ """
+ if not meta or not isinstance(meta, dict):
+ return ''
+ import json
+ from modules.infotext import quote
+ lower = {k.lower(): (k, v) for k, v in meta.items()}
+
+ def lookup(*keys):
+ for k in keys:
+ if k.lower() in lower:
+ return lower[k.lower()][1]
+ return None
+
+ prompt = lookup('prompt') or ''
+ negative = lookup('negativePrompt', 'negative_prompt', 'Negative prompt') or ''
+ pairs = []
+ mapping = [
+ (('steps',), 'Steps'),
+ (('sampler',), 'Sampler'),
+ (('cfgScale', 'cfg_scale', 'CFG scale'), 'CFG scale'),
+ (('seed',), 'Seed'),
+ (('Size',), 'Size'),
+ (('Model',), 'Model'),
+ (('Model hash', 'modelHash'), 'Model hash'),
+ (('clipSkip', 'clip_skip', 'Clip skip'), 'Clip skip'),
+ (('denoisingStrength', 'Denoising strength'), 'Denoising strength'),
+ ]
+ consumed = {'prompt', 'negativeprompt', 'negative_prompt'}
+ for src_keys, out_key in mapping:
+ v = lookup(*src_keys)
+ if v is None or v == '':
+ continue
+ pairs.append(f'{out_key}: {quote(v)}')
+ for k in src_keys:
+ consumed.add(k.lower())
+ for k, v in meta.items():
+ if k.lower() in consumed:
+ continue
+ if k.lower() in NOISE_META_KEYS:
+ continue
+ if k in ('resources', 'civitaiResources'):
+ try:
+ pairs.append(f'Civitai resources: {quote(json.dumps(v, separators=(",", ":")))}')
+ except Exception:
+ pass
+ continue
+ if v is None or v == '':
+ continue
+ quoted = quote(v)
+ if len(str(quoted)) > PASSTHROUGH_VALUE_LIMIT:
+ continue
+ pairs.append(f'{k}: {quoted}')
+ lines = [str(prompt).strip()]
+ if negative:
+ lines.append(f'Negative prompt: {negative}')
+ if pairs:
+ lines.append(', '.join(pairs))
+ return '\n'.join(lines)
+
+
+def fit_parameters_for_exif(parameters: str, limit: int = 30000) -> str:
+ """Trim parameters string to fit JPEG/WEBP EXIF UserComment.
+
+ JPEG's APP1 segment caps at 64KB; UserComment is UTF-16-LE encoded so
+ each char takes 2 bytes. A 30000-char limit keeps the encoded payload
+ around 60KB with headroom for piexif overhead. Drops the
+ `Civitai resources` field first (typical bloat source) and truncates as
+ last resort. PNG callers don't need this since `tEXt` chunks are
+ unbounded.
+ """
+ if len(parameters) <= limit:
+ return parameters
+ lines = parameters.split('\n')
+ if len(lines) >= 3:
+ parts = [p for p in lines[2].split(', ') if not p.startswith('Civitai resources:')]
+ lines[2] = ', '.join(parts)
+ parameters = '\n'.join(lines)
+ if len(parameters) <= limit:
+ return parameters
+ return parameters[:max(0, limit - 32)].rstrip() + '\n[truncated]'
+
+
+def embed_preview_parameters(preview_file: str, parameters: str) -> bool:
+ """Embed parameters string into preview image at `preview_file`.
+
+ PNG -> `tEXt` chunk with key `parameters`. JPEG/WEBP -> EXIF
+ `UserComment` via piexif. RGBA is converted to RGB before JPEG save.
+ Other extensions: no-op. Returns success.
+
+ Writes are atomic (temp file + os.replace). On success, any co-located
+ ` .thumb.jpg` is removed so the lazy thumb generator re-emits it
+ carrying the new params. Embed failures are swallowed and logged at
+ debug; the source file is removed only when PIL cannot read it back,
+ so the caller can re-download a fresh copy.
+ """
+ if not preview_file or not parameters:
+ return False
+ ext = os.path.splitext(preview_file)[1].lower()
+ if ext not in ('.png', '.jpg', '.jpeg', '.webp'):
+ return False
+ tmp_file = preview_file + '.embed.tmp'
+ try:
+ from PIL import Image
+ img = Image.open(preview_file)
+ img.load()
+ try:
+ if ext == '.png':
+ from PIL import PngImagePlugin
+ pnginfo = PngImagePlugin.PngInfo()
+ pnginfo.add_text('parameters', parameters)
+ img.save(tmp_file, format='PNG', pnginfo=pnginfo)
+ else:
+ import piexif
+ import piexif.helper
+ payload = fit_parameters_for_exif(parameters)
+ exif_bytes = piexif.dump({'Exif': {piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(payload, encoding='unicode')}})
+ if ext in ('.jpg', '.jpeg'):
+ if img.mode != 'RGB':
+ img = img.convert('RGB')
+ img.save(tmp_file, format='JPEG', quality=95, exif=exif_bytes)
+ else:
+ img.save(tmp_file, format='WEBP', quality=95, exif=exif_bytes)
+ finally:
+ img.close()
+ os.replace(tmp_file, preview_file)
+ thumb_base = os.path.splitext(preview_file)[0]
+ if thumb_base.endswith('.preview'):
+ thumb_base = thumb_base[:-len('.preview')]
+ thumb_file = thumb_base + '.thumb.jpg'
+ if os.path.exists(thumb_file):
+ try:
+ os.remove(thumb_file)
+ log.debug(f'CivitAI thumb invalidated: file="{thumb_file}"')
+ except Exception:
+ pass
+ return True
+ except Exception as e:
+ try:
+ if os.path.exists(tmp_file):
+ os.remove(tmp_file)
+ except Exception:
+ pass
+ try:
+ from PIL import Image as PILImage
+ with PILImage.open(preview_file) as probe:
+ probe.verify()
+ except Exception:
+ try:
+ os.remove(preview_file)
+ log.warning(f'CivitAI preview removing invalid: image={preview_file}')
+ except Exception:
+ pass
+ log.debug(f'CivitAI preview embed failed: file="{preview_file}" {e}')
+ return False
+
+
+def resolve_preview_file(item: dict) -> str | None:
+ """Resolve the actual on-disk preview file for a network item.
+
+ `item['local_preview']` is the aspirational save path (e.g.
+ ` .`), not necessarily the existing file. This
+ helper extracts the real path from `item['preview']` URL (set by
+ `ExtraNetworksPage.link_preview`) and falls back to scanning common
+ extensions at the model base. Returns None when nothing on-disk
+ matches.
+ """
+ preview_url = item.get('preview') or ''
+ if preview_url and 'missing.png' not in preview_url:
+ try:
+ import urllib.parse
+ parsed = urllib.parse.urlparse(preview_url)
+ fn = urllib.parse.parse_qs(parsed.query).get('filename', [None])[0]
+ if fn:
+ fn = urllib.parse.unquote(fn)
+ if os.path.isfile(fn):
+ return fn
+ except Exception:
+ pass
+ filename = item.get('filename')
+ if not filename:
+ return None
+ return find_ui_preview_file(filename)
+
+
+def preview_has_parameters(preview_file: str) -> bool:
+ """Check whether `preview_file` carries a non-empty embedded parameters string.
+
+ Mirrors `modules/image/metadata.py:read_info_from_image`: PNG
+ `image.info["parameters"]` or JPEG/WEBP EXIF `UserComment`. EXIF
+ UserComment is decoded via piexif.helper to verify a non-empty body;
+ the 8-byte `b"UNICODE\\x00"` header is present in any UserComment
+ even when its body is empty.
+ """
+ if not preview_file or not os.path.exists(preview_file):
+ return False
+ try:
+ from PIL import Image
+ img = Image.open(preview_file)
+ try:
+ info = img.info or {}
+ for key in ('parameters', 'UserComment'):
+ value = info.get(key)
+ if value and str(value).strip():
+ return True
+ exif = info.get('exif')
+ if exif:
+ import piexif
+ import piexif.helper
+ try:
+ parsed = piexif.load(exif)
+ raw = parsed.get('Exif', {}).get(piexif.ExifIFD.UserComment)
+ if raw:
+ try:
+ decoded = piexif.helper.UserComment.load(raw)
+ except Exception:
+ decoded = ''
+ if decoded and decoded.strip():
+ return True
+ except Exception:
+ pass
+ finally:
+ img.close()
+ except Exception:
+ pass
+ return False
+
+
+VIDEO_PREVIEW_EXTENSIONS = ('.mp4', '.webm')
+UI_PREVIEW_EXTS = ('jpg', 'jpeg', 'png', 'webp')
+UI_PREVIEW_MIDS = ('.thumb.', '.', '.preview.')
+
+
+def find_ui_preview_file(model_path: str) -> str | None:
+ """Return the preview file the modernUI surfaces for `model_path`.
+
+ Iteration mirrors `ExtraNetworksPage.find_preview` at
+ `modules/ui_extra_networks.py:488` so that backfill embeds into the
+ same file the UI reads.
+ """
+ base = os.path.splitext(model_path)[0]
+ for ext in UI_PREVIEW_EXTS:
+ for mid in UI_PREVIEW_MIDS:
+ candidate = f'{base}{mid}{ext}'
+ if os.path.isfile(candidate):
+ return candidate
+ return None
+
+
+def backfill_preview_parameters(model_path: str, preview_url: str, meta: dict | None) -> bool:
+ """Embed Civitai meta into an existing preview file when it lacks parameters.
+
+ Used by the rescan path to retroactively populate preview metadata
+ without re-downloading bytes. The embed target is the file
+ `find_ui_preview_file` surfaces, which matches what the modernUI
+ displays. For video previews the embed target is the extracted
+ ` .thumb.jpg` frame instead of the unembeddable video file.
+ Returns True only if a new chunk was written; False for no-op (file
+ missing, no meta, already populated, embed failed).
+ """
+ if not meta or not model_path or not preview_url:
+ return False
+ ext = os.path.splitext(preview_url)[1].lower()
+ base = os.path.splitext(model_path)[0]
+ if ext in VIDEO_PREVIEW_EXTENSIONS:
+ if not os.path.exists(base + ext):
+ return False
+ preview_file = base + '.thumb.jpg'
+ if not os.path.exists(preview_file):
+ return False
+ else:
+ preview_file = find_ui_preview_file(model_path)
+ if not preview_file:
+ return False
+ if preview_has_parameters(preview_file):
+ return False
+ parameters = civitai_meta_to_parameters(meta)
+ if not parameters:
+ return False
+ if embed_preview_parameters(preview_file, parameters):
+ log.info(f'CivitAI preview backfill: file="{preview_file}"')
+ return True
+ return False
+
+
# ---- Legacy compatibility functions ----
def download_civit_meta(model_path: str, model_id):
@@ -361,12 +664,12 @@ def download_civit_meta(model_path: str, model_id):
return r.status_code, '', ''
-def download_civit_preview(model_path: str, preview_url: str):
+def download_civit_preview(model_path: str, preview_url: str, meta: dict | None = None):
if model_path is None:
return 500, '', ''
ext = os.path.splitext(preview_url)[1]
preview_file = os.path.splitext(model_path)[0] + ext
- is_video = preview_file.lower().endswith('.mp4')
+ is_video = preview_file.lower().endswith(VIDEO_PREVIEW_EXTENSIONS)
is_json = preview_file.lower().endswith('.json')
if is_json:
log.warning(f'CivitAI download: url="{preview_url}" skip json')
@@ -389,11 +692,27 @@ def download_civit_preview(model_path: str, preview_url: str):
if is_video:
from modules.civitai.video_helper import save_video_frame
save_video_frame(preview_file)
+ if meta:
+ thumb_file = os.path.splitext(preview_file)[0] + '.thumb.jpg'
+ if os.path.exists(thumb_file):
+ try:
+ parameters = civitai_meta_to_parameters(meta)
+ if parameters and embed_preview_parameters(thumb_file, parameters):
+ log.debug(f'CivitAI preview embed: file="{thumb_file}"')
+ except Exception as e:
+ log.debug(f'CivitAI preview embed skipped: file="{thumb_file}" {e}')
else:
from PIL import Image
img = Image.open(preview_file)
log.info(f'CivitAI download: url={preview_url} file="{preview_file}" size={total_size} image={img.size}')
img.close()
+ if meta:
+ try:
+ parameters = civitai_meta_to_parameters(meta)
+ if parameters and embed_preview_parameters(preview_file, parameters):
+ log.debug(f'CivitAI preview embed: file="{preview_file}"')
+ except Exception as e:
+ log.debug(f'CivitAI preview embed skipped: file="{preview_file}" {e}')
except Exception as e:
log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}')
shared.state.end(jobid)
diff --git a/modules/civitai/metadata_civitai.py b/modules/civitai/metadata_civitai.py
index 626e22967..88f50a1da 100644
--- a/modules/civitai/metadata_civitai.py
+++ b/modules/civitai/metadata_civitai.py
@@ -93,7 +93,7 @@ def civit_update_metadata(raw: bool = False):
model.latest_name = f.get('name', '')
if model.vername == model.latest:
model.status = 'Latest version'
- elif any(map(lambda v: v in model.latest_hashes, all_hashes)): # pylint: disable=cell-var-from-loop # noqa: C417
+ elif any(map(lambda v: v in model.latest_hashes, all_hashes)): # pylint: disable=cell-var-from-loop
model.status = 'Update downloaded'
else:
model.status = 'Update available'
@@ -104,7 +104,7 @@ def civit_update_metadata(raw: bool = False):
def atomic_civit_search_metadata(item, results):
- from modules.civitai.download_civitai import download_civit_preview, download_civit_meta
+ from modules.civitai.download_civitai import download_civit_preview, download_civit_meta, backfill_preview_parameters, preview_has_parameters, resolve_preview_file
if item is None:
return
try:
@@ -112,7 +112,12 @@ def atomic_civit_search_metadata(item, results):
except Exception:
return
has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0
- if ('missing.png' in item['preview'] or not has_meta) and os.path.isfile(item['filename']):
+ needs_backfill = False
+ if has_meta and 'missing.png' not in item.get('preview', ''):
+ actual_preview = resolve_preview_file(item)
+ if actual_preview and not preview_has_parameters(actual_preview):
+ needs_backfill = True
+ if ('missing.png' in item['preview'] or not has_meta or needs_backfill) and os.path.isfile(item['filename']):
sha = item.get('hash', None)
found = False
result = {
@@ -136,11 +141,15 @@ def atomic_civit_search_metadata(item, results):
results.append(dict(result))
for img in version.images:
if img.url:
- code, size, note = download_civit_preview(item['filename'], img.url)
+ code, size, note = download_civit_preview(item['filename'], img.url, meta=img.meta)
if code == 200:
results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'})
found = True
break
+ if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta):
+ results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'})
+ found = True
+ break
else:
result['code'] = 404
time.sleep(0.25) # rate limiting
@@ -157,11 +166,15 @@ def atomic_civit_search_metadata(item, results):
results.append(dict(result))
for img in version.images:
if img.url:
- code, size, note = download_civit_preview(item['filename'], img.url)
+ code, size, note = download_civit_preview(item['filename'], img.url, meta=img.meta)
if code == 200:
results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'})
found = True
break
+ if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta):
+ results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'})
+ found = True
+ break
else:
result['code'] = 404
time.sleep(0.25) # rate limiting
diff --git a/modules/control/proc/canny.py b/modules/control/proc/canny.py
index 1e4bb3176..ad0113a92 100644
--- a/modules/control/proc/canny.py
+++ b/modules/control/proc/canny.py
@@ -7,7 +7,7 @@ from modules.control.util import HWC3, resize_image
class CannyDetector:
def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, image_resolution=512, output_type=None, **kwargs):
if "img" in kwargs:
- warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+ warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning, stacklevel=2)
input_image = kwargs.pop("img")
if input_image is None:
raise ValueError("input_image must be defined.")
diff --git a/modules/control/proc/depth_pro/__init__.py b/modules/control/proc/depth_pro/__init__.py
index e9bd20793..f0f74852e 100644
--- a/modules/control/proc/depth_pro/__init__.py
+++ b/modules/control/proc/depth_pro/__init__.py
@@ -16,7 +16,7 @@ class DepthProDetector:
self.processor = processor
@classmethod
- def from_pretrained(cls, pretrained_model_or_path: str = "apple/DepthPro-hf", cache_dir: str = None, local_files_only = False) -> "DepthProDetector":
+ def from_pretrained(cls, pretrained_model_or_path: str = "apple/DepthPro-hf", cache_dir: str | None = None, local_files_only = False) -> "DepthProDetector":
from transformers import AutoImageProcessor, DepthProForDepthEstimation
processor = AutoImageProcessor.from_pretrained(pretrained_model_or_path, cache_dir=cache_dir, local_files_only=local_files_only)
diff --git a/modules/control/proc/edge.py b/modules/control/proc/edge.py
index f91ab83bb..9de932172 100644
--- a/modules/control/proc/edge.py
+++ b/modules/control/proc/edge.py
@@ -33,7 +33,7 @@ class EdgeDetector:
params.PFmode = pf
ed.setParams(params)
if "img" in kwargs:
- warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+ warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning, stacklevel=2)
input_image = kwargs.pop("img")
if input_image is None:
raise ValueError("input_image must be defined.")
diff --git a/modules/control/proc/leres/leres/network_auxi.py b/modules/control/proc/leres/leres/network_auxi.py
index 34007c9c9..5e9688832 100644
--- a/modules/control/proc/leres/leres/network_auxi.py
+++ b/modules/control/proc/leres/leres/network_auxi.py
@@ -384,7 +384,7 @@ class SenceUnderstand(nn.Module):
self.initial_params()
def forward(self, x):
- n, c, h, w = x.size()
+ n, _c, h, w = x.size()
x = self.conv1(x)
x = self.pool(x)
x = x.view(n, -1)
diff --git a/modules/control/proc/leres/pix2pix/options/base_options.py b/modules/control/proc/leres/pix2pix/options/base_options.py
index 533a1e88a..63acea872 100644
--- a/modules/control/proc/leres/pix2pix/options/base_options.py
+++ b/modules/control/proc/leres/pix2pix/options/base_options.py
@@ -1,8 +1,8 @@
import argparse
import os
-from ...pix2pix.util import util
+from proc.leres.pix2pix.util import util
# import torch
-from ...pix2pix import models
+from proc.leres.pix2pix import models
# import pix2pix.data
import numpy as np
diff --git a/modules/control/proc/marigold/marigold_pipeline.py b/modules/control/proc/marigold/marigold_pipeline.py
index 768c67684..4ab8c0495 100644
--- a/modules/control/proc/marigold/marigold_pipeline.py
+++ b/modules/control/proc/marigold/marigold_pipeline.py
@@ -113,7 +113,7 @@ class MarigoldPipeline(DiffusionPipeline):
batch_size: int = 0,
color_map: str = "Spectral",
show_progress_bar: bool = True,
- ensemble_kwargs: Dict = None,
+ ensemble_kwargs: Dict | None = None,
) -> MarigoldDepthOutput:
"""
Function invoked when calling the pipeline.
diff --git a/modules/control/proc/marigold/util/ensemble.py b/modules/control/proc/marigold/util/ensemble.py
index 710db1cc2..1139b5554 100644
--- a/modules/control/proc/marigold/util/ensemble.py
+++ b/modules/control/proc/marigold/util/ensemble.py
@@ -43,7 +43,7 @@ def ensemble_depths(
max_iter: int = 2,
tol: float = 1e-3,
reduction: str = "median",
- max_res: int = None,
+ max_res: int | None = None,
):
"""
To ensemble multiple affine-invariant depth images (up to scale and shift),
diff --git a/modules/control/proc/marigold/util/seed_all.py b/modules/control/proc/marigold/util/seed_all.py
index b09006c9b..b3cdcaae5 100644
--- a/modules/control/proc/marigold/util/seed_all.py
+++ b/modules/control/proc/marigold/util/seed_all.py
@@ -28,6 +28,6 @@ def seed_all(seed: int = 0):
Set random seeds of all components.
"""
random.seed(seed)
- np.random.seed(seed) # noqa
+ np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
diff --git a/modules/control/proc/midas/midas/vit.py b/modules/control/proc/midas/midas/vit.py
index f268a9fc4..0f2aae284 100644
--- a/modules/control/proc/midas/midas/vit.py
+++ b/modules/control/proc/midas/midas/vit.py
@@ -54,7 +54,7 @@ class Transpose(nn.Module):
def forward_vit(pretrained, x):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
pretrained.model.forward_flex(x)
@@ -115,7 +115,7 @@ def _resize_pos_embed(self, posemb, gs_h, gs_w):
def forward_flex(self, x):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
pos_embed = self._resize_pos_embed(
self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
diff --git a/modules/control/proc/midas/utils.py b/modules/control/proc/midas/utils.py
index 9a9d3b5b6..f3de55e64 100644
--- a/modules/control/proc/midas/utils.py
+++ b/modules/control/proc/midas/utils.py
@@ -75,7 +75,7 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (
- len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
+ len(image.shape) == 2 or (len(image.shape) == 3 and image.shape[2] == 1)
): # greyscale
color = False
else:
@@ -86,7 +86,7 @@ def write_pfm(path, image, scale=1):
endian = image.dtype.byteorder
- if endian == "<" or endian == "=" and sys.byteorder == "little":
+ if endian == "<" or (endian == "=" and sys.byteorder == "little"):
scale = -scale
file.write("%f\n".encode() % scale)
diff --git a/modules/control/proc/mlsd/utils.py b/modules/control/proc/mlsd/utils.py
index ca8034370..7915c3d48 100644
--- a/modules/control/proc/mlsd/utils.py
+++ b/modules/control/proc/mlsd/utils.py
@@ -22,7 +22,7 @@ def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
center: tpMap[1, 0, :, :]
displacement: tpMap[1, 1:5, :, :]
'''
- b, c, h, w = tpMap.shape
+ b, _c, _h, w = tpMap.shape
assert b==1, 'only support bsize==1'
displacement = tpMap[:, 1:5, :, :][0]
center = tpMap[:, 0, :, :]
@@ -471,9 +471,9 @@ def pred_squares(image,
square[end_idx]
# check whether outside or inside
- start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
+ _start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
connect_idx)
- end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
+ _end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
cover += dist_segments[connect_idx] + start_cover_param * start_min + end_cover_param * end_min
perimeter += dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min
diff --git a/modules/control/proc/openpose/__init__.py b/modules/control/proc/openpose/__init__.py
index 746351718..4d4c72911 100644
--- a/modules/control/proc/openpose/__init__.py
+++ b/modules/control/proc/openpose/__init__.py
@@ -194,14 +194,14 @@ class OpenposeDetector:
def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=False, include_face=False, hand_and_face=None, output_type="pil", **kwargs):
self.to(devices.device)
if hand_and_face is not None:
- warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning)
+ warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning, stacklevel=2)
include_hand = hand_and_face
include_face = hand_and_face
if "return_pil" in kwargs:
- warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+ warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning, stacklevel=2)
output_type = "pil" if kwargs["return_pil"] else "np"
if type(output_type) is bool:
- warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+ warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions", stacklevel=2)
if output_type:
output_type = "pil"
if not isinstance(input_image, np.ndarray):
diff --git a/modules/control/proc/openpose/face.py b/modules/control/proc/openpose/face.py
index e8e34451c..43e7d04cb 100644
--- a/modules/control/proc/openpose/face.py
+++ b/modules/control/proc/openpose/face.py
@@ -328,7 +328,7 @@ class Face(object):
def __call__(self, face_img):
device = next(iter(self.model.parameters())).device
- H, W, C = face_img.shape
+ H, W, _C = face_img.shape
w_size = 384
x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5
diff --git a/modules/control/proc/openpose/hand.py b/modules/control/proc/openpose/hand.py
index 78e00213e..576ad7172 100644
--- a/modules/control/proc/openpose/hand.py
+++ b/modules/control/proc/openpose/hand.py
@@ -32,7 +32,7 @@ class Hand(object):
wsize = 128
heatmap_avg = np.zeros((wsize, wsize, 22))
- Hr, Wr, Cr = oriImgRaw.shape
+ Hr, Wr, _Cr = oriImgRaw.shape
oriImg = cv2.GaussianBlur(oriImgRaw, (0, 0), 0.8)
diff --git a/modules/control/proc/segment_anything/__init__.py b/modules/control/proc/segment_anything/__init__.py
index 121421c18..a86c7d6d7 100644
--- a/modules/control/proc/segment_anything/__init__.py
+++ b/modules/control/proc/segment_anything/__init__.py
@@ -53,7 +53,7 @@ class SamDetector:
def __call__(self, input_image: Union[np.ndarray, Image.Image]=None, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs) -> Image.Image:
if "image" in kwargs:
- warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+ warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning, stacklevel=2)
input_image = kwargs.pop("image")
if input_image is None:
raise ValueError("input_image must be defined.")
diff --git a/modules/control/proc/segment_anything/modeling/sam.py b/modules/control/proc/segment_anything/modeling/sam.py
index 614fd7483..d99cd3387 100644
--- a/modules/control/proc/segment_anything/modeling/sam.py
+++ b/modules/control/proc/segment_anything/modeling/sam.py
@@ -25,8 +25,8 @@ class Sam(nn.Module):
image_encoder: Union[ImageEncoderViT, TinyViT],
prompt_encoder: PromptEncoder,
mask_decoder: MaskDecoder,
- pixel_mean: List[float] = None,
- pixel_std: List[float] = None,
+ pixel_mean: List[float] | None = None,
+ pixel_std: List[float] | None = None,
) -> None:
"""
SAM predicts object masks from an image and input prompts.
diff --git a/modules/control/proc/segment_anything/modeling/transformer.py b/modules/control/proc/segment_anything/modeling/transformer.py
index 28fafea52..5d6155002 100644
--- a/modules/control/proc/segment_anything/modeling/transformer.py
+++ b/modules/control/proc/segment_anything/modeling/transformer.py
@@ -79,7 +79,7 @@ class TwoWayTransformer(nn.Module):
torch.Tensor: the processed image_embedding
"""
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
- bs, c, h, w = image_embedding.shape
+ _bs, _c, _h, _w = image_embedding.shape
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
image_pe = image_pe.flatten(2).permute(0, 2, 1)
diff --git a/modules/control/proc/segment_anything/utils/onnx.py b/modules/control/proc/segment_anything/utils/onnx.py
index 103867faf..4eaf0b89b 100644
--- a/modules/control/proc/segment_anything/utils/onnx.py
+++ b/modules/control/proc/segment_anything/utils/onnx.py
@@ -10,7 +10,7 @@ from torch.nn import functional as F
from typing import Tuple
-from ..modeling import Sam
+from proc.segment_anything.modeling import Sam
from .amg import calculate_stability_score
diff --git a/modules/control/proc/shuffle.py b/modules/control/proc/shuffle.py
index 3ee285857..01da6719c 100644
--- a/modules/control/proc/shuffle.py
+++ b/modules/control/proc/shuffle.py
@@ -10,10 +10,10 @@ from modules.control.util import HWC3, img2mask, make_noise_disk, resize_image
class ContentShuffleDetector:
def __call__(self, input_image, h=None, w=None, f=None, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
if "return_pil" in kwargs:
- warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+ warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning, stacklevel=2)
output_type = "pil" if kwargs["return_pil"] else "np"
if type(output_type) is bool:
- warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+ warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions", stacklevel=2)
if output_type:
output_type = "pil"
@@ -49,7 +49,7 @@ class ContentShuffleDetector:
class ColorShuffleDetector:
def __call__(self, img):
H, W, C = img.shape
- F = np.random.randint(64, 384) # noqa
+ F = np.random.randint(64, 384)
A = make_noise_disk(H, W, 3, F)
B = make_noise_disk(H, W, 3, F)
C = (A + B) / 2.0
@@ -82,11 +82,11 @@ class DownSampleDetector:
def __call__(self, img, level=3, k=16.0):
h = img.astype(np.float32)
for _ in range(level):
- h += np.random.normal(loc=0.0, scale=k, size=h.shape) # noqa
+ h += np.random.normal(loc=0.0, scale=k, size=h.shape)
h = cv2.pyrDown(h)
for _ in range(level):
h = cv2.pyrUp(h)
- h += np.random.normal(loc=0.0, scale=k, size=h.shape) # noqa
+ h += np.random.normal(loc=0.0, scale=k, size=h.shape)
return h.clip(0, 255).astype(np.uint8)
diff --git a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/beit.py b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/beit.py
index ab7458704..cd835ebc1 100644
--- a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/beit.py
+++ b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/beit.py
@@ -66,7 +66,7 @@ def attention_forward(self, x, resolution, shared_rel_pos_bias: Optional[torch.T
"""
Modification of timm.models.beit.py: Attention.forward to support arbitrary window sizes.
"""
- B, N, C = x.shape
+ B, N, _C = x.shape
qkv_bias = torch.cat((self.q_bias, self.k_bias, self.v_bias)) if self.q_bias is not None else None
qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
diff --git a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/utils.py b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/utils.py
index bed17f97d..a58fa876b 100644
--- a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/utils.py
+++ b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/utils.py
@@ -81,7 +81,7 @@ def forward_default(pretrained, x, function_name="forward_features"):
def forward_adapted_unflatten(pretrained, x, function_name="forward_features"):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
exec(f"glob = pretrained.model.{function_name}(x)")
diff --git a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/vit.py b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/vit.py
index 71e864cdf..ce8ea0176 100644
--- a/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/vit.py
+++ b/modules/control/proc/zoe/zoedepth/models/base_models/midas_repo/midas/backbones/vit.py
@@ -31,7 +31,7 @@ def _resize_pos_embed(self, posemb, gs_h, gs_w):
def forward_flex(self, x):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
pos_embed = self._resize_pos_embed(
self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
diff --git a/modules/control/proc/zoe/zoedepth/models/layers/attractor.py b/modules/control/proc/zoe/zoedepth/models/layers/attractor.py
index c2fe653ed..5e58ee348 100644
--- a/modules/control/proc/zoe/zoedepth/models/layers/attractor.py
+++ b/modules/control/proc/zoe/zoedepth/models/layers/attractor.py
@@ -100,7 +100,7 @@ class AttractorLayer(nn.Module):
A = self._net(x)
eps = 1e-3
A = A + eps
- n, c, h, w = A.shape
+ n, _c, h, w = A.shape
A = A.view(n, self.n_attractors, 2, h, w)
A_normed = A / A.sum(dim=2, keepdim=True) # n, a, 2, h, w
A_normed = A[:, :, 0, ...] # n, na, h, w
@@ -177,7 +177,7 @@ class AttractorLayerUnnormed(nn.Module):
x = x + prev_b_embedding
A = self._net(x)
- n, c, h, w = A.shape
+ _n, _c, h, w = A.shape
b_prev = nn.functional.interpolate(
b_prev, (h, w), mode='bilinear', align_corners=True)
diff --git a/modules/control/proc/zoe/zoedepth/models/layers/localbins_layers.py b/modules/control/proc/zoe/zoedepth/models/layers/localbins_layers.py
index b70ae562e..9af4dc463 100644
--- a/modules/control/proc/zoe/zoedepth/models/layers/localbins_layers.py
+++ b/modules/control/proc/zoe/zoedepth/models/layers/localbins_layers.py
@@ -146,7 +146,7 @@ class LinearSplitter(nn.Module):
S = self._net(x)
eps = 1e-3
S = S + eps
- n, c, h, w = S.shape
+ n, _c, h, w = S.shape
S = S.view(n, self.prev_nbins, self.split_factor, h, w)
S_normed = S / S.sum(dim=2, keepdim=True) # fractional splits
diff --git a/modules/control/proc/zoe/zoedepth/models/zoedepth/zoedepth_v1.py b/modules/control/proc/zoe/zoedepth/models/zoedepth/zoedepth_v1.py
index 1705442c6..b10afa59f 100644
--- a/modules/control/proc/zoe/zoedepth/models/zoedepth/zoedepth_v1.py
+++ b/modules/control/proc/zoe/zoedepth/models/zoedepth/zoedepth_v1.py
@@ -26,13 +26,12 @@ import itertools
import torch
import torch.nn as nn
-from ..depth_model import DepthModel
-from ..base_models.midas import MidasCore
-from ..layers.attractor import AttractorLayer, AttractorLayerUnnormed
-from ..layers.dist_layers import ConditionalLogBinomial
-from ..layers.localbins_layers import (Projector, SeedBinRegressor,
- SeedBinRegressorUnnormed)
-from ..model_io import load_state_from_resource
+from proc.zoe.zoedepth.models.depth_model import DepthModel
+from proc.zoe.zoedepth.models.base_models.midas import MidasCore
+from proc.zoe.zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
+from proc.zoe.zoedepth.models.layers.dist_layers import ConditionalLogBinomial
+from proc.zoe.zoedepth.models.layers.localbins_layers import Projector, SeedBinRegressor, SeedBinRegressorUnnormed
+from proc.zoe.zoedepth.models.model_io import load_state_from_resource
class ZoeDepth(DepthModel):
@@ -139,7 +138,7 @@ class ZoeDepth(DepthModel):
- probs (torch.Tensor): Output probability distribution of shape (B, n_bins, H, W). Present only if return_probs is True
"""
- b, c, h, w = x.shape
+ b, _c, h, w = x.shape
# print("input shape ", x.shape)
self.orig_input_width = w
self.orig_input_height = h
diff --git a/modules/control/proc/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py b/modules/control/proc/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py
index 889b1e282..52f427add 100644
--- a/modules/control/proc/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py
+++ b/modules/control/proc/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py
@@ -27,14 +27,13 @@ import itertools
import torch
import torch.nn as nn
-from ..depth_model import DepthModel
-from ..base_models.midas import MidasCore
-from ..layers.attractor import AttractorLayer, AttractorLayerUnnormed
-from ..layers.dist_layers import ConditionalLogBinomial
-from ..layers.localbins_layers import (Projector, SeedBinRegressor,
- SeedBinRegressorUnnormed)
-from ..layers.patch_transformer import PatchTransformerEncoder
-from ..model_io import load_state_from_resource
+from proc.zoe.zoedepth.models.depth_model import DepthModel
+from proc.zoe.zoedepth.models.base_models.midas import MidasCore
+from proc.zoe.zoedepth.models.layers.attractor import AttractorLayer, AttractorLayerUnnormed
+from proc.zoe.zoedepth.models.layers.dist_layers import ConditionalLogBinomial
+from proc.zoe.zoedepth.models.layers.localbins_layers import Projector, SeedBinRegressor, SeedBinRegressorUnnormed
+from proc.zoe.zoedepth.models.layers.patch_transformer import PatchTransformerEncoder
+from proc.zoe.zoedepth.models.model_io import load_state_from_resource
class ZoeDepthNK(DepthModel):
def __init__(self, core, bin_conf, bin_centers_type="softplus", bin_embedding_dim=128,
@@ -173,10 +172,10 @@ class ZoeDepthNK(DepthModel):
- "bin_centers": Bin centers of shape (B, N, H, W). Present only if return_final_centers is True
- "probs": Bin probabilities of shape (B, N, H, W). Present only if return_probs is True
"""
- b, c, h, w = x.shape
+ b, _c, h, w = x.shape
self.orig_input_width = w
self.orig_input_height = h
- rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
+ _rel_depth, out = self.core(x, denorm=denorm, return_rel_depth=True)
outconv_activation = out[0]
btlnck = out[1]
diff --git a/modules/control/processors.py b/modules/control/processors.py
index bdf5e53c7..1f03f5594 100644
--- a/modules/control/processors.py
+++ b/modules/control/processors.py
@@ -386,7 +386,9 @@ class Processor:
def preview(self):
import modules.ui_control_helpers as helpers
input_image = helpers.input_source
+ if input_image is None:
+ return []
if isinstance(input_image, list):
input_image = input_image[0]
debug('Control process preview')
- return self.__call__(input_image)
+ return [self.__call__(input_image)]
diff --git a/modules/control/run.py b/modules/control/run.py
index c201367f5..0ba71021d 100644
--- a/modules/control/run.py
+++ b/modules/control/run.py
@@ -274,10 +274,73 @@ def init_units(units: list[unit.Unit]):
u.process.override = u.override
+def control_process(p: StableDiffusionProcessingControl,
+ input_script_args: list | None = None,
+ override_script_name: str | None = None,
+ override_script_args: list | None = None,
+ input_image: Image.Image = None, # only used for tiling, otherwise processor.preprocess_image set p params
+ ):
+ debug_log(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
+ if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE and hasattr(pipe, 'vae'): # force vae back to gpu if not in txt2img mode
+ sd_models.move_model(pipe.vae, devices.device)
+
+ # what are we doing?
+ if 'control' in p.ops:
+ p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_control_samples)
+ elif 'img2img' in p.ops:
+ p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_img2img_samples)
+ elif 'txt2img' in p.ops:
+ p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_txt2img_samples)
+ else: # fallback to txt2img
+ p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_txt2img_samples)
+
+ # init scripts
+ p.scripts = scripts_manager.scripts_control
+ p.script_args = input_script_args or []
+ if len(p.script_args) == 0:
+ if not p.scripts:
+ p.scripts.initialize_scripts(False)
+ p.script_args = script.init_default_script_args(p.scripts)
+
+ # init override scripts
+ if override_script_name and override_script_args and len(override_script_name) > 0:
+ selectable_scripts, selectable_script_idx = script.get_selectable_script(override_script_name, p.scripts)
+ if selectable_scripts:
+ for idx in range(len(override_script_args)):
+ p.script_args[selectable_scripts.args_from + idx] = override_script_args[idx]
+ p.script_args[0] = selectable_script_idx + 1
+
+ # actual processing
+ script_run = False
+ processed: processing.Processed = None
+ if p.is_tile:
+ processed: processing.Processed = tile.run_tiling(p, input_image)
+ if processed is None and p.scripts is not None:
+ processed = p.scripts.run(p, *p.script_args)
+ if processed is None:
+ processed: processing.Processed = processing.process_images(p) # run actual pipeline
+ else:
+ script_run = True
+
+ # postprocessing
+ if p.scripts is not None:
+ processed = p.scripts.after(p, processed, *p.script_args)
+
+ output = None
+ info = None
+ if processed is not None and processed.images is not None:
+ output = processed.images
+ info = [processed.infotext(p, i) for i in range(len(output))]
+
+ # output = pipe(**vars(p)).images # debug: direct pipe exec call
+ return output, info, script_run
+
+
def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
units: list[unit.Unit] | None = None, inputs: list[Image.Image] | None = None, inits: list[Image.Image] | None = None, mask: Image.Image = None, unit_type: str | None = None, is_generator: bool = True,
input_type: int = 0,
- prompt: str = '', negative_prompt: str = '', styles: list[str] | None = None,
+ prompt: str = '', negative_prompt: str = '',
+ styles: list[str] | None = None,
steps: int = 20, sampler_index: int | None = None,
seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1,
guidance_name: str = 'Default', guidance_scale: float = 6.0, guidance_rescale: float = 0.0, guidance_start: float = 0.0, guidance_stop: float = 1.0,
@@ -296,7 +359,9 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
resize_mode_before: int = 0, resize_name_before: str = 'None', resize_context_before: str = 'None', width_before: int = 512, height_before: int = 512, scale_by_before: float = 1.0, selected_scale_tab_before: int = 0,
resize_mode_after: int = 0, resize_name_after: str = 'None', resize_context_after: str = 'None', width_after: int = 0, height_after: int = 0, scale_by_after: float = 1.0, selected_scale_tab_after: int = 0,
resize_mode_mask: int = 0, resize_name_mask: str = 'None', resize_context_mask: str = 'None', width_mask: int = 0, height_mask: int = 0, scale_by_mask: float = 1.0, selected_scale_tab_mask: int = 0,
- denoising_strength: float = 0.3, batch_count: int = 1, batch_size: int = 1,
+ denoising_strength: float = 0.3,
+ skip_processing: bool = False,
+ batch_count: int = 1, batch_size: int = 1,
enable_hr: bool = False, hr_sampler_index: int | None = None, hr_denoising_strength: float = 0.0, hr_resize_mode: int = 0, hr_resize_context: str = 'None', hr_upscaler: str | None = None, hr_force: bool = False, hr_second_pass_steps: int = 20,
hr_scale: float = 1.0, hr_resize_x: int = 0, hr_resize_y: int = 0, refiner_steps: int = 5, refiner_start: float = 0.0, refiner_prompt: str = '', refiner_negative: str = '',
video_skip_frames: int = 0, video_type: str = 'None', video_duration: float = 2.0, video_loop: bool = False, video_pad: int = 0, video_interpolate: int = 0,
@@ -394,6 +459,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
seed_resize_from_h = seed_resize_from_h,
seed_resize_from_w = seed_resize_from_w,
denoising_strength = denoising_strength,
+ skip_processing = skip_processing,
# modular guidance
guidance_name = guidance_name,
guidance_scale = guidance_scale,
@@ -633,17 +699,20 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
processed_image = None
if frame is not None:
inputs = [Image.fromarray(frame)] # cv2 to pil
- for i, input_image in enumerate(inputs):
- if input_image is not None:
- p.ops.append('img2img')
+ for i, input_image in enumerate(inputs): # loop per-input, but with early-break
if pipe is None: # pipe may have been reset externally
if video is None:
break # non-video: pipeline was consumed, no need to re-process remaining inputs
pipe = set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, active_units, control_conditioning, control_guidance_start, control_guidance_end, inits)
debug_log(f'Control pipeline reinit: class={pipe.__class__.__name__}')
+
pipe.restore_pipeline = restore_pipeline
shared.sd_model.restore_pipeline = restore_pipeline
debug_log(f'Control Control image: {i + 1} of {len(inputs)}')
+
+ if input_image is not None:
+ p.ops.append('img2img')
+
if shared.state.skipped:
shared.state.skipped = False
continue
@@ -652,6 +721,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
if is_generator:
yield terminate('Interrupted')
return terminate('Interrupted')
+
# get input
if isinstance(input_image, str) and os.path.exists(input_image):
try:
@@ -664,7 +734,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
debug_log('Control Init image: same as control')
init_image = input_image
elif inits is None or len(inits) == 0:
- debug_log('Control Init image: none')
+ debug_log('Control init image: none')
init_image = None
elif len(inits) > i and isinstance(inits[i], str):
debug_log(f'Control: init image: {inits[i]}')
@@ -683,9 +753,25 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
continue
index += 1
- processed_image, blended_image = preprocess_image(p, pipe, input_image, init_image, mask, input_type, unit_type, active_process, active_model, selected_models, has_models, active_units)
- if is_generator:
- yield (None, blended_image, '') # result is control_output, proces_output
+ if getattr(pipe, 'skip_processing', False) or getattr(p, 'skip_processing', False):
+ p.init_images = inputs
+ p.extra_generation_params['Process'] = False
+ else:
+ processed_image, blended_image = preprocess_image(p,
+ pipe,
+ input_image,
+ init_image,
+ mask,
+ input_type,
+ unit_type,
+ active_process,
+ active_model,
+ selected_models,
+ has_models,
+ active_units,
+ )
+ if is_generator:
+ yield (None, blended_image, '') # result is control_output, proces_output
# final check
if has_models:
@@ -701,62 +787,18 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
if unit_type == 'lite':
instance.apply(selected_models, processed_image, control_conditioning)
- # what are we doing?
- if 'control' in p.ops:
- p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_control_samples)
- elif 'img2img' in p.ops:
- p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_img2img_samples)
- elif 'txt2img' in p.ops:
- p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_txt2img_samples)
- else: # fallback to txt2img
- p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_txt2img_samples)
-
# pipeline
output = None
script_run = False
- if pipe is not None: # run new pipeline
- debug_log(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
- if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE: # force vae back to gpu if not in txt2img mode
- sd_models.move_model(pipe.vae, devices.device)
-
- # init scripts
- p.scripts = scripts_manager.scripts_control
- p.script_args = input_script_args or []
- if len(p.script_args) == 0:
- if not p.scripts:
- p.scripts.initialize_scripts(False)
- p.script_args = script.init_default_script_args(p.scripts)
-
- # init override scripts
- if override_script_name and override_script_args and len(override_script_name) > 0:
- selectable_scripts, selectable_script_idx = script.get_selectable_script(override_script_name, p.scripts)
- if selectable_scripts:
- for idx in range(len(override_script_args)):
- p.script_args[selectable_scripts.args_from + idx] = override_script_args[idx]
- p.script_args[0] = selectable_script_idx + 1
-
- # actual processing
- processed: processing.Processed = None
- if p.is_tile:
- processed: processing.Processed = tile.run_tiling(p, input_image)
- if processed is None and p.scripts is not None:
- processed = p.scripts.run(p, *p.script_args)
- if processed is None:
- processed: processing.Processed = processing.process_images(p) # run actual pipeline
- else:
- script_run = True
-
- # postprocessing
- if p.scripts is not None:
- processed = p.scripts.after(p, processed, *p.script_args)
- output = None
- if processed is not None and processed.images is not None:
- output = processed.images
- info_txt = [processed.infotext(p, i) for i in range(len(output))]
-
- # output = pipe(**vars(p)).images # alternative direct pipe exec call
- else: # blend all processed images and return
+ if pipe is None: # blend all processed images and return
output = processed_image
+ else: # run new pipeline
+ output, info_txt, script_run = control_process(p,
+ input_script_args,
+ override_script_name,
+ override_script_args,
+ input_image,
+ )
# outputs
output = output or []
@@ -810,7 +852,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
html_txt = f'Ready {image_txt}
' if image_txt != '' else ''
if len(info_txt) > 0:
html_txt = html_txt + infotext_to_html(info_txt[0])
+ result = (output_images, blended_image, html_txt, output_filename)
if is_generator:
- yield (output_images, blended_image, html_txt, output_filename)
- else:
- return (output_images, blended_image, html_txt, output_filename)
+ yield result
+ return result
diff --git a/modules/control/test.py b/modules/control/test.py
index 29653e41a..6627c7d04 100644
--- a/modules/control/test.py
+++ b/modules/control/test.py
@@ -39,7 +39,7 @@ def test_processors(image):
draw.text((10, 10), processor_id, (0,0,0), font=font)
draw.text((8, 8), processor_id, (255,255,255), font=font)
res.append(output)
- yield output, None, None, res
+ yield output, None, res
rows = round(math.sqrt(len(res)))
cols = math.ceil(len(res) / rows)
w, h = 256, 256
@@ -52,8 +52,7 @@ def test_processors(image):
thumb = image.copy().convert('RGB')
thumb.thumbnail((w, h), Image.Resampling.HAMMING)
grid.paste(thumb, box=(x, y))
- yield None, grid, None, res
- return None, grid, None, res # preview_process, output_image, output_video, output_gallery
+ yield None, grid, res
def test_controlnets(prompt, negative, image):
@@ -90,7 +89,7 @@ def test_controlnets(prompt, negative, image):
draw.text((10, 10), model_id, (0,0,0), font=font)
draw.text((8, 8), model_id, (255,255,255), font=font)
res.append(output)
- yield output, None, None, res
+ yield output, None, res
rows = round(math.sqrt(len(res)))
cols = math.ceil(len(res) / rows)
w, h = 256, 256
@@ -103,8 +102,7 @@ def test_controlnets(prompt, negative, image):
thumb = image.copy().convert('RGB')
thumb.thumbnail((w, h), Image.Resampling.HAMMING)
grid.paste(thumb, box=(x, y))
- yield None, grid, None, res
- return None, grid, None, res # preview_process, output_image, output_video, output_gallery
+ yield grid, None, res
def test_adapters(prompt, negative, image):
@@ -142,7 +140,7 @@ def test_adapters(prompt, negative, image):
draw.text((10, 10), model_id, (0,0,0), font=font)
draw.text((8, 8), model_id, (255,255,255), font=font)
res.append(output)
- yield output, None, None, res
+ yield output, None, res
rows = round(math.sqrt(len(res)))
cols = math.ceil(len(res) / rows)
w, h = 256, 256
@@ -155,8 +153,7 @@ def test_adapters(prompt, negative, image):
thumb = image.copy().convert('RGB')
thumb.thumbnail((w, h), Image.Resampling.HAMMING)
grid.paste(thumb, box=(x, y))
- yield None, grid, None, res
- return None, grid, None, res # preview_process, output_image, output_video, output_gallery
+ yield grid, None, res
def test_xs(prompt, negative, image):
@@ -193,7 +190,7 @@ def test_xs(prompt, negative, image):
draw.text((10, 10), model_id, (0,0,0), font=font)
draw.text((8, 8), model_id, (255,255,255), font=font)
res.append(output)
- yield output, None, None, res
+ yield output, None, res
rows = round(math.sqrt(len(res)))
cols = math.ceil(len(res) / rows)
w, h = 256, 256
@@ -206,8 +203,7 @@ def test_xs(prompt, negative, image):
thumb = image.copy().convert('RGB')
thumb.thumbnail((w, h), Image.Resampling.HAMMING)
grid.paste(thumb, box=(x, y))
- yield None, grid, None, res
- return None, grid, None, res # preview_process, output_image, output_video, output_gallery
+ yield grid, None, res
def test_lite(prompt, negative, image):
@@ -245,7 +241,7 @@ def test_lite(prompt, negative, image):
draw.text((10, 10), model_id, (0,0,0), font=font)
draw.text((8, 8), model_id, (255,255,255), font=font)
res.append(output)
- yield output, None, None, res
+ yield output, None, res
rows = round(math.sqrt(len(res)))
cols = math.ceil(len(res) / rows)
w, h = 256, 256
@@ -258,5 +254,5 @@ def test_lite(prompt, negative, image):
thumb = image.copy().convert('RGB')
thumb.thumbnail((w, h), Image.Resampling.HAMMING)
grid.paste(thumb, box=(x, y))
- yield None, grid, None, res
- return None, grid, None, res # preview_process, output_image, output_video, output_gallery
+ yield grid, None, res
+ return grid, None, res # preview_process, output_image, output_video, output_gallery
diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py
index 5bc231f56..8efbce5a5 100644
--- a/modules/control/units/controlnet.py
+++ b/modules/control/units/controlnet.py
@@ -5,7 +5,7 @@ from typing import Union
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline, StableDiffusion3Pipeline, ControlNetModel
from modules.control.units import detect
from modules.shared import log, opts, cmd_opts, state, listdir
-from modules import errors, sd_models, devices, model_quant
+from modules import errors, sd_models, devices, model_quant # pylint: disable=unused-import
from modules.processing import StableDiffusionProcessingControl
@@ -163,7 +163,7 @@ def find_models():
find_models()
-def api_list_models(model_type: str = None):
+def api_list_models(model_type: str | None = None):
import modules.shared
model_type = model_type or modules.shared.sd_model_type
model_list = []
@@ -215,7 +215,7 @@ def list_models(refresh=False):
class ControlNet():
- def __init__(self, model_id: str = None, device = None, dtype = None, load_config = None):
+ def __init__(self, model_id: str | None = None, device = None, dtype = None, load_config = None):
self.model: ControlNetModel = None
self.model_id: str = model_id
self.device = device
@@ -311,34 +311,34 @@ class ControlNet():
self.load_config['original_config_file '] = config_path
self.model = cls.from_single_file(model_path, config=config, **self.load_config)
- def load(self, model_id: str = None, force: bool = False) -> str:
+ def load(self, model_id: str | None = None, force: bool = False) -> str:
with load_lock:
try:
t0 = time.time()
model_id = model_id or self.model_id
if model_id is None or model_id == 'None':
self.reset()
- return
+ return ''
if model_id not in all_models:
log.error(f'Control {what}: id="{model_id}" available={list(all_models)} unknown model')
- return
+ return ''
model_path = all_models[model_id]
if model_path == '':
- return
+ return ''
if model_path is None:
log.error(f'Control {what} model load: id="{model_id}" unknown model id')
- return
+ return ''
if 'lora' in model_id.lower():
self.model = model_path
- return
+ return ''
if model_id == self.model_id and not force:
# log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
- return
+ return ''
log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
cls, config = self.get_class(model_id)
if cls is None:
log.error(f'Control {what} model load: id="{model_id}" unknown base model')
- return
+ return ''
self.reset()
jobid = state.begin(f'Load {what}')
if model_path.endswith('.safetensors'):
@@ -363,7 +363,7 @@ class ControlNet():
if debug:
errors.display(e, 'Control')
if self.model is None:
- return
+ return ''
if not cmd_opts.lowvram: # lowvram will cause unet<->controlnet to ping-pong but saves more memory
self.model.offload_never = True
if self.dtype is not None:
diff --git a/modules/control/units/lite.py b/modules/control/units/lite.py
index 107ebb0a0..c09a48f5e 100644
--- a/modules/control/units/lite.py
+++ b/modules/control/units/lite.py
@@ -63,7 +63,7 @@ def list_models(refresh=False):
class ControlLLLite():
- def __init__(self, model_id: str = None, device = None, dtype = None, load_config = None):
+ def __init__(self, model_id: str | None = None, device = None, dtype = None, load_config = None):
self.model: ControlNetLLLite = None
self.model_id: str = model_id
self.device = device
@@ -83,26 +83,25 @@ class ControlLLLite():
self.model = None
self.model_id = None
- def load(self, model_id: str = None, force: bool = True) -> str:
+ def load(self, model_id: str | None = None, force: bool = True) -> str:
with load_lock:
try:
t0 = time.time()
model_id = model_id or self.model_id
if model_id is None or model_id == 'None':
self.reset()
- return
+ return ''
if model_id not in all_models:
log.error(f'Control {what} unknown model: id="{model_id}" available={list(all_models)}')
- return
+ return ''
model_path = all_models[model_id]
if model_path == '':
- return
+ return ''
if model_path is None:
log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
- return
+ return ''
if model_id == self.model_id and not force:
- # log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
- return
+ return ''
log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}" {self.load_config}')
if model_path.endswith('.safetensors'):
self.model = ControlNetLLLite(model_path)
diff --git a/modules/control/units/t2iadapter.py b/modules/control/units/t2iadapter.py
index 5bdf99fb8..697a948fe 100644
--- a/modules/control/units/t2iadapter.py
+++ b/modules/control/units/t2iadapter.py
@@ -71,7 +71,7 @@ class AdapterModel(T2IAdapter):
class Adapter():
- def __init__(self, model_id: str = None, device = None, dtype = None, load_config = None):
+ def __init__(self, model_id: str | None = None, device = None, dtype = None, load_config = None):
self.model: AdapterModel = None
self.model_id: str = model_id
self.device = device
@@ -91,17 +91,17 @@ class Adapter():
self.model = None
self.model_id = None
- def load(self, model_id: str = None, force: bool = True) -> str:
+ def load(self, model_id: str | None = None, force: bool = True) -> str:
with load_lock:
try:
t0 = time.time()
model_id = model_id or self.model_id
if model_id is None or model_id == 'None':
self.reset()
- return
+ return ''
if model_id not in all_models:
log.error(f'Control {what} unknown model: id="{model_id}" available={list(all_models)}')
- return
+ return ''
model_path, model_args = all_models[model_id]
self.load_config.update(model_args)
from modules.shared import opts
@@ -113,10 +113,9 @@ class Adapter():
os.unsetenv('HF_HUB_OFFLINE')
if model_path is None:
log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
- return
+ return ''
if model_id == self.model_id and not force:
- # log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
- return
+ return ''
log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
if model_path.endswith('.pth') or model_path.endswith('.pt') or model_path.endswith('.safetensors') or model_path.endswith('.bin'):
from huggingface_hub import hf_hub_download
diff --git a/modules/control/units/xs.py b/modules/control/units/xs.py
index 2d56fd7ff..bfcc8d44e 100644
--- a/modules/control/units/xs.py
+++ b/modules/control/units/xs.py
@@ -59,7 +59,7 @@ def list_models(refresh=False):
class ControlNetXS():
- def __init__(self, model_id: str = None, device = None, dtype = None, load_config = None):
+ def __init__(self, model_id: str | None = None, device = None, dtype = None, load_config = None):
self.model: ControlNetXSModel = None
self.model_id: str = model_id
self.device = device
@@ -79,26 +79,25 @@ class ControlNetXS():
self.model = None
self.model_id = None
- def load(self, model_id: str = None, time_embedding_mix: float = 0.0, force: bool = True) -> str:
+ def load(self, model_id: str | None = None, time_embedding_mix: float = 0.0, force: bool = True) -> str:
with load_lock:
try:
t0 = time.time()
model_id = model_id or self.model_id
if model_id is None or model_id == 'None':
self.reset()
- return
+ return ''
if model_id not in all_models:
log.error(f'Control {what} unknown model: id="{model_id}" available={list(all_models)}')
- return
+ return ''
model_path = all_models[model_id]
if model_path == '':
- return
+ return ''
if model_path is None:
log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
- return
+ return ''
if model_id == self.model_id and not force:
- # log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
- return
+ return ''
self.load_config['time_embedding_mix'] = time_embedding_mix
if opts.offline_mode:
self.load_config["local_files_only"] = True
diff --git a/modules/control/units/xs_model.py b/modules/control/units/xs_model.py
index f4866b58d..9040d24a7 100644
--- a/modules/control/units/xs_model.py
+++ b/modules/control/units/xs_model.py
@@ -64,9 +64,9 @@ class ControlNetXSOutput(BaseOutput):
class ControlNetConditioningEmbedding(nn.Module):
"""
Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN
- [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized
- training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the
- convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides
+ [11] to convert the entire dataset of 512 x 512 images into smaller 64 x 64 “latent images” for stabilized
+ training. This requires ControlNets to convert image-based conditions to 64 x 64 feature space to match the
+ convolution size. We use a tiny network E(·) of four convolution layers with 4 x 4 kernels and 2 x 2 strides
(activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full
model) to encode image-space conditions ... into feature maps ..."
"""
@@ -657,7 +657,7 @@ class ControlNetXSModel(ModelMixin, ConfigMixin):
if base_model.config.addition_embed_type == "text":
aug_emb = base_model.add_embedding(encoder_hidden_states)
elif base_model.config.addition_embed_type == "text_image":
- raise NotImplementedError()
+ raise NotImplementedError
elif base_model.config.addition_embed_type == "text_time":
# SDXL - style
if "text_embeds" not in added_cond_kwargs:
@@ -676,9 +676,9 @@ class ControlNetXSModel(ModelMixin, ConfigMixin):
add_embeds = add_embeds.to(temb.dtype)
aug_emb = base_model.add_embedding(add_embeds)
elif base_model.config.addition_embed_type == "image":
- raise NotImplementedError()
+ raise NotImplementedError
elif base_model.config.addition_embed_type == "image_hint":
- raise NotImplementedError()
+ raise NotImplementedError
temb = temb + aug_emb if aug_emb is not None else temb
diff --git a/modules/control/units/xs_pipe.py b/modules/control/units/xs_pipe.py
index f178c11b1..878bb5c66 100644
--- a/modules/control/units/xs_pipe.py
+++ b/modules/control/units/xs_pipe.py
@@ -518,23 +518,23 @@ class StableDiffusionXLControlNetXSPipeline(
)
if (
isinstance(self.controlnet, ControlNetXSModel)
- or is_compiled
- and isinstance(self.controlnet._orig_mod, ControlNetXSModel)
+ or (is_compiled
+ and isinstance(self.controlnet._orig_mod, ControlNetXSModel))
):
self.check_image(image, prompt, prompt_embeds)
else:
- assert False
+ raise AssertionError
# Check `controlnet_conditioning_scale`
if (
isinstance(self.controlnet, ControlNetXSModel)
- or is_compiled
- and isinstance(self.controlnet._orig_mod, ControlNetXSModel)
+ or (is_compiled
+ and isinstance(self.controlnet._orig_mod, ControlNetXSModel))
):
if not isinstance(controlnet_conditioning_scale, float):
raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
else:
- assert False
+ raise AssertionError
start, end = control_guidance_start, control_guidance_end
if start >= end:
@@ -700,7 +700,7 @@ class StableDiffusionXLControlNetXSPipeline(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
image: PipelineImageInput = None,
height: Optional[int] = None,
@@ -725,9 +725,9 @@ class StableDiffusionXLControlNetXSPipeline(
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
control_guidance_start: float = 0.0,
control_guidance_end: float = 1.0,
- original_size: Tuple[int, int] = None,
+ original_size: Tuple[int, int] | None = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ target_size: Tuple[int, int] | None = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
@@ -928,7 +928,7 @@ class StableDiffusionXLControlNetXSPipeline(
)
height, width = image.shape[-2:]
else:
- assert False
+ raise AssertionError
# 5. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
@@ -1521,23 +1521,23 @@ class StableDiffusionControlNetXSPipeline(
)
if (
isinstance(self.controlnet, ControlNetXSModel)
- or is_compiled
- and isinstance(self.controlnet._orig_mod, ControlNetXSModel)
+ or (is_compiled
+ and isinstance(self.controlnet._orig_mod, ControlNetXSModel))
):
self.check_image(image, prompt, prompt_embeds)
else:
- assert False
+ raise AssertionError
# Check `controlnet_conditioning_scale`
if (
isinstance(self.controlnet, ControlNetXSModel)
- or is_compiled
- and isinstance(self.controlnet._orig_mod, ControlNetXSModel)
+ or (is_compiled
+ and isinstance(self.controlnet._orig_mod, ControlNetXSModel))
):
if not isinstance(controlnet_conditioning_scale, float):
raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
else:
- assert False
+ raise AssertionError
start, end = control_guidance_start, control_guidance_end
if start >= end:
@@ -1664,7 +1664,7 @@ class StableDiffusionControlNetXSPipeline(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
image: PipelineImageInput = None,
height: Optional[int] = None,
width: Optional[int] = None,
@@ -1831,7 +1831,7 @@ class StableDiffusionControlNetXSPipeline(
)
height, width = image.shape[-2:]
else:
- assert False
+ raise AssertionError
# 5. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
diff --git a/modules/extra_networks.py b/modules/extra_networks.py
index 26e9d7829..c281c521f 100644
--- a/modules/extra_networks.py
+++ b/modules/extra_networks.py
@@ -37,6 +37,15 @@ class ExtraNetworkParams:
else:
self.positional.append(item)
+ def __eq__(self, other):
+ if not isinstance(other, ExtraNetworkParams):
+ return False
+ same = self.positional == other.positional and self.named == other.named
+ return same
+
+ def __repr__(self):
+ return f"NetworkParams(positional={self.positional} named={self.named})"
+
class ExtraNetwork:
def __init__(self, name):
@@ -66,14 +75,6 @@ class ExtraNetwork:
raise NotImplementedError
-def is_stepwise(en_obj):
- all_args = []
- for en in en_obj:
- all_args.extend(en.positional[1:])
- all_args.extend(en.named.values())
- return any([len(str(x).split("@")) > 1 for x in all_args]) # noqa C419 # pylint: disable=use-a-generator
-
-
def activate(p, extra_network_data=None, step=0, include=None, exclude=None):
"""call activate for extra networks in extra_network_data in specified order, then call activate for all remaining registered networks with an empty argument list"""
if exclude is None:
@@ -83,16 +84,6 @@ def activate(p, extra_network_data=None, step=0, include=None, exclude=None):
if p.disable_extra_networks:
return
extra_network_data = extra_network_data or p.network_data
- # if extra_network_data is None or len(extra_network_data) == 0:
- # return
- stepwise = False
- for extra_network_args in extra_network_data.values():
- stepwise = stepwise or is_stepwise(extra_network_args)
- functional = shared.opts.lora_functional
- if shared.opts.lora_force_diffusers and stepwise:
- log.warning("Network load: type=LoRA method=composable loader=diffusers not compatible")
- stepwise = False
- shared.opts.data['lora_functional'] = stepwise or functional
for extra_network_name, extra_network_args in extra_network_data.items():
extra_network = extra_network_registry.get(extra_network_name, None)
@@ -122,9 +113,6 @@ def activate(p, extra_network_data=None, step=0, include=None, exclude=None):
errors.display(e, f"Activating network: type={extra_network_name}")
p.network_data = extra_network_data
- if stepwise:
- p.stepwise_lora = True
- shared.opts.data['lora_functional'] = functional
def deactivate(p, extra_network_data=None, force=None):
@@ -134,8 +122,7 @@ def deactivate(p, extra_network_data=None, force=None):
if force is None:
force = shared.opts.lora_force_reload
extra_network_data = extra_network_data or p.network_data
- # if extra_network_data is None or len(extra_network_data) == 0:
- # return
+
for extra_network_name in extra_network_data:
extra_network = extra_network_registry.get(extra_network_name, None)
if extra_network is None:
diff --git a/modules/face/faceid.py b/modules/face/faceid.py
index f3cd6f783..b2263d840 100644
--- a/modules/face/faceid.py
+++ b/modules/face/faceid.py
@@ -214,7 +214,7 @@ def face_id(
p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n+1) * p.batch_size]
p.seeds = p.all_seeds[n * p.batch_size:(n+1) * p.batch_size]
p.subseeds = p.all_subseeds[n * p.batch_size:(n+1) * p.batch_size]
- p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts)
+ p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data)
extra_networks.activate(p, p.network_data)
ip_model_dict.update({
diff --git a/modules/flash_attn_triton_amd/utils.py b/modules/flash_attn_triton_amd/utils.py
index 8b1640ff2..503298f09 100644
--- a/modules/flash_attn_triton_amd/utils.py
+++ b/modules/flash_attn_triton_amd/utils.py
@@ -110,7 +110,7 @@ class MetaData():
def check_args(self, q, k, v, o):
assert q.dim() == k.dim() and q.dim() == v.dim()
- batch, nheads_q, nheads_k, head_size, _, _ = get_shapes_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k)
+ _batch, nheads_q, nheads_k, _head_size, _, _ = get_shapes_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k)
if self.varlen:
assert q.dim() == 3
assert self.cu_seqlens_q is not None
@@ -310,7 +310,7 @@ def input_helper(
return qkv, do, metadata
else:
- assert False, f"Unsupported packing mode: {packing}"
+ raise AssertionError(f"Unsupported packing mode: {packing}")
# -------------------------------
# Alibi
@@ -361,7 +361,7 @@ def get_shape_from_layout(
elif layout == 'bshd':
batch, max_seqlen_final, num_heads, head_dim = x.shape
elif layout == 'thd':
- total_seqlen, num_heads, head_dim = x.shape
+ _total_seqlen, num_heads, head_dim = x.shape
if cu_seqlens is None:
raise ValueError("cu_seqlens must be provided for varlen (thd) layout")
if max_seqlen is None:
@@ -369,7 +369,7 @@ def get_shape_from_layout(
batch, max_seqlen_final, num_heads, head_dim = len(cu_seqlens) - 1, max_seqlen, num_heads, head_dim
else:
- assert False, "Got unsupported layout."
+ raise AssertionError("Got unsupported layout.")
return batch, max_seqlen_final, num_heads, head_dim
@@ -392,7 +392,7 @@ def get_stride_from_layout(x: torch.Tensor, layout:Literal["bshd", "bhsd", "thd"
elif layout == 'bshd':
strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3))
else:
- assert False, 'Got unsupported layout.'
+ raise AssertionError('Got unsupported layout.')
return strides
def get_shape_and_strides_from_layout(x: torch.Tensor, layout: Literal["bshd", "bhsd", "thd"], cu_seqlens: Optional[torch.Tensor] = None, max_seqlen: Optional[int] = None):
diff --git a/modules/framepack/framepack_vae.py b/modules/framepack/framepack_vae.py
index 7596f2f9a..b06ea3fa4 100644
--- a/modules/framepack/framepack_vae.py
+++ b/modules/framepack/framepack_vae.py
@@ -45,7 +45,7 @@ def vae_decode_tiny(latents):
global taesd # pylint: disable=global-statement
if taesd is None:
from modules.vae import sd_vae_taesd
- taesd, _variant = sd_vae_taesd.get_model(variant='TAE HunyuanVideo')
+ taesd, _variant = sd_vae_taesd.load_model(variant='TAE HunyuanVideo')
log.debug(f'Video VAE: type=Tiny cls={taesd.__class__.__name__} latents={latents.shape}')
with devices.inference_context():
taesd = taesd.to(device=devices.device, dtype=devices.dtype)
diff --git a/modules/framepack/framepack_wrappers.py b/modules/framepack/framepack_wrappers.py
index bdff734ae..b9a5e5592 100644
--- a/modules/framepack/framepack_wrappers.py
+++ b/modules/framepack/framepack_wrappers.py
@@ -70,7 +70,7 @@ def prepare_prompts(p, init_image, prompt:str, section_prompt:str, num_sections:
p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
shared.prompt_styles.apply_styles_to_extra(p)
- p.prompts, p.network_data = extra_networks.parse_prompts([p.prompt])
+ p.prompts, p.network_data = extra_networks.parse_prompts([p.prompt], p.network_data)
extra_networks.activate(p)
prompt = p.prompts[0]
generated_prompts = [''] * num_sections
diff --git a/modules/framepack/pipeline/hunyuan_video_packed.py b/modules/framepack/pipeline/hunyuan_video_packed.py
index a852fbc10..a5f171b5c 100644
--- a/modules/framepack/pipeline/hunyuan_video_packed.py
+++ b/modules/framepack/pipeline/hunyuan_video_packed.py
@@ -832,7 +832,7 @@ class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterM
clean_latents_4x=None, clean_latent_4x_indices=None
):
hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents)
- B, C, T, H, W = hidden_states.shape
+ B, _C, T, H, W = hidden_states.shape
if latent_indices is None:
latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1)
@@ -897,7 +897,7 @@ class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterM
if attention_kwargs is None:
attention_kwargs = {}
- batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ batch_size, _num_channels, num_frames, height, width = hidden_states.shape
p, p_t = self.config['patch_size'], self.config['patch_size_t']
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p
diff --git a/modules/gr_hijack.py b/modules/gr_hijack.py
index 0fe7eb24e..082ace269 100644
--- a/modules/gr_hijack.py
+++ b/modules/gr_hijack.py
@@ -153,6 +153,7 @@ def patch_gradio():
orig_cancel_tasks = gradio.utils.cancel_tasks
orig_restore_session_state = gradio.route_utils.restore_session_state
orig_call_prediction = gradio.queueing.Queue.call_prediction
+ orig_blocks_preprocess_data = gradio.blocks.Blocks.preprocess_data
async def wrap_cancel_tasks(task_ids: set[str]):
log.error(f'Gradio cancel: task={task_ids}')
@@ -200,9 +201,17 @@ def patch_gradio():
log.error(f"Gradio queue: events={len(events)} batch={batch} error: {e}")
raise
+ def wrap_blocks_preprocess_data(self, fn_index: int, inputs: list, state: dict):
+ try:
+ return orig_blocks_preprocess_data(self, fn_index, inputs, state)
+ except Exception as e:
+ log.error(f"Gradio preprocess: {e}")
+ raise
+
gradio.queueing.Queue.call_prediction = wrap_call_prediction
gradio.route_utils.restore_session_state = wrap_restore_session_state
gradio.utils.cancel_tasks = wrap_cancel_tasks
+ gradio.blocks.Blocks.preprocess_data = wrap_blocks_preprocess_data
def patch_gradio_future():
diff --git a/modules/hidiffusion/hidiffusion.py b/modules/hidiffusion/hidiffusion.py
index b0ddf3b58..fd2c4b408 100644
--- a/modules/hidiffusion/hidiffusion.py
+++ b/modules/hidiffusion/hidiffusion.py
@@ -107,7 +107,7 @@ def make_diffusers_transformer_block(block_class: Type[torch.nn.Module]) -> Type
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
timestep: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
+ cross_attention_kwargs: Dict[str, Any] | None = None,
class_labels: Optional[torch.LongTensor] = None,
added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
) -> torch.FloatTensor:
diff --git a/modules/hidiffusion/hidiffusion_controlnet.py b/modules/hidiffusion/hidiffusion_controlnet.py
index 990ecf8da..ef7ce0457 100644
--- a/modules/hidiffusion/hidiffusion_controlnet.py
+++ b/modules/hidiffusion/hidiffusion_controlnet.py
@@ -272,7 +272,7 @@ def make_diffusers_sdxl_contrtolnet_ppl(block_class):
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
@@ -298,9 +298,9 @@ def make_diffusers_sdxl_contrtolnet_ppl(block_class):
guess_mode: bool = False,
control_guidance_start: Union[float, List[float]] = 0.0,
control_guidance_end: Union[float, List[float]] = 1.0,
- original_size: Tuple[int, int] = None,
+ original_size: Tuple[int, int] | None = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ target_size: Tuple[int, int] | None = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
diff --git a/modules/history.py b/modules/history.py
index 62166b895..f98fea1b2 100644
--- a/modules/history.py
+++ b/modules/history.py
@@ -80,7 +80,6 @@ class History:
break
current_index -= 1
if item.latent is None:
- print('HERE')
return None, -1
log.debug(f'History get: index={current_index} time={item.ts} shape={list(item.latent.shape)} dtype={item.latent.dtype} count={self.count}')
return item.latent.to(devices.device), current_index
diff --git a/modules/image/namegen.py b/modules/image/namegen.py
index be850e6b9..d69e0e08d 100644
--- a/modules/image/namegen.py
+++ b/modules/image/namegen.py
@@ -55,7 +55,7 @@ class FilenameGenerator:
'prompt_hash': lambda self: hashlib.sha256(self.prompt.encode()).hexdigest()[0:8],
'sampler': lambda self: self.p and self.p.sampler_name,
- 'seed': lambda self: (self.seed and str(self.seed)) or '',
+ 'seed': lambda self: (hasattr(self, "seed") and self.seed and str(self.seed)) or '',
'steps': lambda self: self.p and getattr(self.p, 'steps', 0),
'cfg': lambda self: self.p and getattr(self.p, 'cfg_scale', 0),
'pag': lambda self: self.p and getattr(self.p, 'pag_scale', 0),
diff --git a/modules/images.py b/modules/images.py
index d294e02b7..05c33d802 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -3,6 +3,7 @@ from modules.image.save import save_image, sanitize_filename_part
from modules.image.resize import resize_image
from modules.image.namegen import FilenameGenerator, get_next_sequence_number
from modules.image.grid import Grid, image_grid, check_grid_size, get_grid_size, draw_grid_annotations, draw_prompt_matrix, combine_grid, get_font
+from modules.image.util import draw_text, flatten
__all__ = [
'check_grid_size',
@@ -20,4 +21,6 @@ __all__ = [
'save_image',
'get_font',
'get_next_sequence_number',
+ 'draw_text',
+ 'flatten',
]
diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py
index 363f3a991..5c36efd67 100644
--- a/modules/intel/ipex/diffusers.py
+++ b/modules/intel/ipex/diffusers.py
@@ -50,7 +50,7 @@ def hidream_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
omega = 1.0 / (theta**scale)
- batch_size, seq_length = pos.shape
+ batch_size, _seq_length = pos.shape
out = torch.einsum("...n,d->...nd", pos, omega)
cos_out = torch.cos(out)
sin_out = torch.sin(out)
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index e16ee2d89..6fb0714c0 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/modules/intel/ipex/hijacks.py
@@ -52,7 +52,7 @@ def autocast_init(self, device_type=None, dtype=None, enabled=True, cache_enable
original_grad_scaler_init = torch.amp.grad_scaler.GradScaler.__init__
@wraps(torch.amp.grad_scaler.GradScaler.__init__)
-def GradScaler_init(self, device: str = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True):
+def GradScaler_init(self, device: str | None = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True):
if device is None or check_cuda(device):
return original_grad_scaler_init(self, device=return_xpu(device), init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled)
else:
@@ -62,10 +62,11 @@ def GradScaler_init(self, device: str = None, init_scale: float = 2.0**16, growt
original_is_autocast_enabled = torch.is_autocast_enabled
@wraps(torch.is_autocast_enabled)
def torch_is_autocast_enabled(device_type=None):
- if device_type is None or check_cuda(device_type):
- return original_is_autocast_enabled(return_xpu(device_type))
- else:
- return original_is_autocast_enabled(device_type)
+ dev = device_type
+ if dev is None or check_cuda(dev):
+ dev = return_xpu(dev)
+ dev = str(dev) if not isinstance(dev, str) else dev
+ return original_is_autocast_enabled(dev)
original_get_autocast_dtype = torch.get_autocast_dtype
diff --git a/modules/ipadapter.py b/modules/ipadapter.py
index 291ae896a..3b9d7147e 100644
--- a/modules/ipadapter.py
+++ b/modules/ipadapter.py
@@ -201,10 +201,10 @@ def load_image_encoder(pipe: DiffusionPipeline, adapter_names: list[str]):
else:
if clip_subfolder is None:
image_encoder = transformers.CLIPVisionModelWithProjection.from_pretrained(clip_repo, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, use_safetensors=True, **offline_config)
- log.debug(f'IP adapter load: encoder="{clip_repo}" cls={pipe.image_encoder.__class__.__name__}')
+ log.debug(f'IP adapter load: encoder="{clip_repo}" cls={image_encoder.__class__.__name__}')
else:
image_encoder = transformers.CLIPVisionModelWithProjection.from_pretrained(clip_repo, subfolder=clip_subfolder, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, use_safetensors=True, **offline_config)
- log.debug(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" cls={pipe.image_encoder.__class__.__name__}')
+ log.debug(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" cls={image_encoder.__class__.__name__}')
sd_models.clear_caches()
image_encoder = model_quant.do_post_load_quant(image_encoder, allow=True)
if hasattr(pipe, 'register_modules'):
@@ -212,12 +212,16 @@ def load_image_encoder(pipe: DiffusionPipeline, adapter_names: list[str]):
else:
pipe.image_encoder = image_encoder
clip_loaded = f'{clip_repo}/{clip_subfolder}'
+ pipe = sd_models.apply_balanced_offload(pipe, force=True)
except Exception as e:
log.error(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" {e}')
errors.display(e, 'IP adapter: type=encoder')
return False
shared.state.end(jobid)
sd_models.move_model(pipe.image_encoder, devices.device)
+ if hasattr(pipe.unet, 'balanced_offload_device_map') and pipe.unet.balanced_offload_device_map.get('encoder_hid_proj', None) is None:
+ # image encoder patches unet, but diffusers creates full device map only during model load so if module is loaded later it will be missing
+ pipe.unet.balanced_offload_device_map['encoder_hid_proj'] = 0
return True
@@ -235,13 +239,14 @@ def load_feature_extractor(pipe):
pipe.register_modules(feature_extractor=feature_extractor)
else:
pipe.feature_extractor = feature_extractor
- sd_models.apply_balanced_offload(pipe.feature_extractor)
+ pipe = sd_models.apply_balanced_offload(pipe, force=True, silent=True)
log.debug(f'IP adapter load: extractor={pipe.feature_extractor.__class__.__name__}')
except Exception as e:
log.error(f'IP adapter load: extractor {e}')
errors.display(e, 'IP adapter: type=extractor')
return False
shared.state.end(jobid)
+ sd_models.move_model(pipe.feature_extractor, devices.device)
return True
diff --git a/modules/logger.py b/modules/logger.py
index 89f5a2277..d912a4cfe 100644
--- a/modules/logger.py
+++ b/modules/logger.py
@@ -62,11 +62,28 @@ def install_traceback(suppress: list | None = None):
-_log_config = {'debug': False, 'trace': False, 'log_filename': None}
+_log_config = {'debug': False, 'trace': False, 'log_filename': None, 'traceback_panel_patched': False}
def setup_logging(debug=None, trace=None, filename=None):
global log_file, console, log_rolled # pylint: disable=global-statement
+ def patch_traceback_panel():
+ # Rich traceback always wraps output in Panel; replace only traceback-local Panel with passthrough.
+ if str_to_bool(os.environ.get("SD_TRACEBOX", False)):
+ return
+ try:
+ import rich.traceback as rich_traceback
+ if _log_config['traceback_panel_patched']:
+ return
+
+ def _panel_passthrough(renderable, *args, **kwargs): # pylint: disable=unused-argument
+ return renderable
+
+ rich_traceback.Panel = _panel_passthrough
+ _log_config['traceback_panel_patched'] = True
+ except Exception:
+ pass
+
if debug is not None:
_log_config['debug'] = debug
if trace is not None:
@@ -177,6 +194,7 @@ def setup_logging(debug=None, trace=None, filename=None):
Padding.__rich_console__ = override_padding
box.ROUNDED = box.SIMPLE
+ patch_traceback_panel()
console = Console(
log_time=True,
log_time_format='%H:%M:%S-%f',
diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py
index 364244b1a..a12452a8c 100644
--- a/modules/lora/extra_networks_lora.py
+++ b/modules/lora/extra_networks_lora.py
@@ -177,11 +177,11 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
def changed(self, requested: list[str], include: list[str] | None = None, exclude: list[str] | None = None) -> bool:
if shared.opts.lora_force_reload:
- debug_log(f'Network check: type=LoRA requested={requested} status=forced')
- return True
+ debug_log(f'Network check: type=LoRA requested={requested} status="forced"')
+ return True, "forced"
sd_model = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
if sd_model is None:
- return False
+ return False, "no model"
if not hasattr(sd_model, 'loaded_loras'):
sd_model.loaded_loras = {}
if include is None or len(include) == 0:
@@ -192,15 +192,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
loaded = sd_model.loaded_loras.get(key, [])
if len(requested) != len(loaded):
sd_model.loaded_loras[key] = requested
- debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status=changed')
- return True
+ debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="num changed"')
+ return True, "num changed"
for req, load in zip(requested, loaded, strict=False):
if req != load:
sd_model.loaded_loras[key] = requested
- debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status=changed')
- return True
- debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status=same')
- return False
+ debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="content changed"')
+ return True, "content changed"
+ debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="same"')
+ return False, "none"
def activate(self, p, params_list, step=0, include=None, exclude=None):
if exclude is None:
@@ -216,15 +216,16 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
self.model = shared.opts.sd_model_checkpoint
names, te_multipliers, unet_multipliers, dyn_dims, lora_modules = parse(p, params_list, step)
requested = self.signature(names, te_multipliers, unet_multipliers)
+ reason = ''
- load_method = lora_overrides.get_method()
+ load_method, load_reason = lora_overrides.get_method()
if debug:
import sys
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
- debug_log(f'Network load: type=LoRA include={include} exclude={exclude} method={load_method} requested={requested} fn={fn}')
+ debug_log(f'Network load: type=LoRA include={include} exclude={exclude} method={load_method} reason={load_reason} requested={requested} fn={fn}')
if load_method == 'diffusers':
- has_changed = self.changed(requested)
+ has_changed, reason = self.changed(requested)
if has_changed:
jobid = shared.state.begin('LoRA')
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, lora_modules) # load only on first call
@@ -237,7 +238,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
else: # native
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
- has_changed = self.changed(requested, include, exclude)
+ has_changed, reason = self.changed(requested, include, exclude)
if has_changed:
jobid = shared.state.begin('LoRA')
if len(l.previously_loaded_networks) > 0:
@@ -254,7 +255,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
prompt(p)
if has_changed and len(include) == 0: # print only once
actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method
- log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} method={actual_method} mode={"fuse" if shared.opts.lora_fuse_native else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary}')
+ log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if shared.opts.lora_fuse_native else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason={reason}')
def deactivate(self, p, force=False):
if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force):
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
index 43543a042..c6b1c1816 100644
--- a/modules/lora/lora_apply.py
+++ b/modules/lora/lora_apply.py
@@ -112,7 +112,7 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
else:
batch_updown = updown.to(devices.device)
if ex_bias is not None:
- if batch_ex_bias:
+ if batch_ex_bias is not None:
batch_ex_bias += ex_bias.to(batch_ex_bias.device)
else:
batch_ex_bias = ex_bias.to(devices.device)
diff --git a/modules/lora/lora_diffusers.py b/modules/lora/lora_diffusers.py
index a2fc0ff27..94ad716c7 100644
--- a/modules/lora/lora_diffusers.py
+++ b/modules/lora/lora_diffusers.py
@@ -50,11 +50,11 @@ def load_per_module(sd_model: diffusers.DiffusionPipeline, filename: str, adapte
return adapter_name
-def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale:float=shared.opts.extra_networks_default_multiplier, lora_module=None) -> network.Network | None:
+def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale:float=shared.opts.extra_networks_default_multiplier, lora_module=None, reason: str = '') -> network.Network | None:
t0 = time.time()
name = name.replace(".", "_")
sd_model: diffusers.DiffusionPipeline = getattr(shared.sd_model, "pipe", shared.sd_model)
- log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
+ log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
if not hasattr(sd_model, 'load_lora_weights'):
log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora')
return None
diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py
index e1536198a..19a55d209 100644
--- a/modules/lora/lora_extract.py
+++ b/modules/lora/lora_extract.py
@@ -238,7 +238,7 @@ def create_ui():
with gr.Tab(label="Extract LoRA"):
with gr.Row():
- gr.HTML(' Extract currently loaded LoRA(s) ')
+ gr.HTML(' Extract currently loaded LoRA(s) ')
with gr.Row():
loaded = gr.Textbox(placeholder="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora_str()}, "lora_extract_refresh")
diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py
index f9cea425d..b3c4b0908 100644
--- a/modules/lora/lora_load.py
+++ b/modules/lora/lora_load.py
@@ -71,10 +71,46 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
if anima_net is not None:
lora_cache[name] = anima_net
return anima_net
+ if shared.sd_model_type == 'ernieimage':
+ from pipelines.ernie import ernie_lora
+ lora_scale = shared.opts.extra_networks_default_multiplier
+ ernie_net = None
+ for try_fn in (ernie_lora.try_load_lora, ernie_lora.try_load_lokr, ernie_lora.try_load_loha, ernie_lora.try_load_oft):
+ sub = try_fn(name, network_on_disk, lora_scale)
+ if sub is None:
+ continue
+ if ernie_net is None:
+ ernie_net = sub
+ else:
+ ernie_net.modules.update(sub.modules)
+ if ernie_net is not None:
+ lora_cache[name] = ernie_net
+ return ernie_net
+ if shared.sd_model_type == 'chroma':
+ from pipelines.chroma import chroma_lora
+ lora_scale = shared.opts.extra_networks_default_multiplier
+ chroma_net = None
+ for try_fn in (chroma_lora.try_load_lora, chroma_lora.try_load_lokr, chroma_lora.try_load_loha, chroma_lora.try_load_oft):
+ sub = try_fn(name, network_on_disk, lora_scale)
+ if sub is None:
+ continue
+ if chroma_net is None:
+ chroma_net = sub
+ else:
+ chroma_net.modules.update(sub.modules)
+ if chroma_net is not None:
+ lora_cache[name] = chroma_net
+ return chroma_net
+ if shared.sd_model_type == 'f2':
+ from pipelines.flux import flux2_lora
+ f2_net = flux2_lora.try_load(name, network_on_disk, shared.opts.extra_networks_default_multiplier)
+ if f2_net is not None:
+ lora_cache[name] = f2_net
+ return f2_net
net = network.Network(name, network_on_disk)
net.mtime = os.path.getmtime(network_on_disk.filename)
state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
- if shared.sd_model_type in ['f1', 'chroma']: # if kohya flux lora, convert state_dict
+ if shared.sd_model_type == 'f1': # if kohya flux lora, convert state_dict
state_dict = lora_convert._convert_kohya_flux_lora_to_diffusers(state_dict) or state_dict # pylint: disable=protected-access
if shared.sd_model_type == 'sd3': # if kohya flux lora, convert state_dict
try:
@@ -277,15 +313,9 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
lora_module = lora_modules[i] if lora_modules and len(lora_modules) > i else None
if recompile_model and shared.compiled_model_state is not None:
shared.compiled_model_state.lora_model.append(f"{name}:{lora_scale}")
- lora_method = lora_overrides.get_method(shorthash)
+ lora_method, lora_method_reason = lora_overrides.get_method(shorthash)
if lora_method == 'diffusers':
- if shared.sd_model_type == 'f2':
- from pipelines.flux import flux2_lora
- net = flux2_lora.try_load_lokr(name, network_on_disk, lora_scale)
- if net is None and not shared.opts.lora_force_diffusers:
- net = flux2_lora.try_load_lora(name, network_on_disk, lora_scale)
- if net is None:
- net = lora_diffusers.load_diffusers(name, network_on_disk, lora_scale, lora_module)
+ net = lora_diffusers.load_diffusers(name, network_on_disk, lora_scale, lora_module, reason=lora_method_reason)
elif lora_method == 'nunchaku':
pass # handled directly from extra_networks_lora.load_nunchaku
else:
@@ -342,8 +372,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
# Also restore backed-up weights when previously active native modules are removed
from modules.lora import networks
native_nets = [net for net in l.loaded_networks if len(net.modules) > 0]
- had_native = len(networks.applied_layers) > 0
- if native_nets or had_native:
+ if native_nets or networks.native_active:
networks.network_activate()
if len(l.loaded_networks) > 0 and l.debug:
diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py
index dc43e2220..efd842c3e 100644
--- a/modules/lora/lora_overrides.py
+++ b/modules/lora/lora_overrides.py
@@ -26,9 +26,11 @@ allow_native = [
'sdxl',
'sd3',
'f1',
+ 'f2',
'chroma',
'zimage',
'anima',
+ 'ernieimage',
]
@@ -42,18 +44,36 @@ fuse_ignore = [
def get_method(shorthash=''):
- use_diffusers = shared.opts.lora_force_diffusers or (shared.sd_model.__class__.__name__ in force_classes_diffusers) or (shared.sd_model_type not in allow_native)
- if len(shorthash) > 4:
- use_diffusers = use_diffusers or any(x.startswith(shorthash) for x in force_hashes_diffusers)
+ """Return ``(method, reason)`` for the active LoRA loading strategy.
+
+ ``method`` is one of ``'native'``, ``'diffusers'``, ``'nunchaku'``.
+ ``reason`` is a short identifier indicating which condition triggered the
+ chosen method, useful for distinguishing user-opt-in from automatic
+ fallback in logs. Reasons:
+
+ - ``'nunchaku-transformer'`` / ``'nunchaku-unet'``: a Nunchaku-quantized
+ component is loaded.
+ - ``'opt-in'``: ``shared.opts.lora_force_diffusers`` is on (settings).
+ - ``'class-forced'``: pipeline class is in ``force_classes_diffusers``.
+ - ``'arch-unsupported'``: ``sd_model_type`` is not in ``allow_native``.
+ - ``'hash-forced'``: file hash is in ``force_hashes_diffusers``.
+ - ``'default'``: native path is the active and unforced choice.
+ """
nunchaku_dit = hasattr(shared.sd_model, 'transformer') and 'Nunchaku' in shared.sd_model.transformer.__class__.__name__
nunchaku_unet = hasattr(shared.sd_model, 'unet') and 'Nunchaku' in shared.sd_model.unet.__class__.__name__
- use_nunchaku = nunchaku_dit or nunchaku_unet
- if use_nunchaku:
- return 'nunchaku'
- elif use_diffusers:
- return 'diffusers'
- else:
- return 'native'
+ if nunchaku_dit:
+ return 'nunchaku', 'nunchaku-transformer'
+ if nunchaku_unet:
+ return 'nunchaku', 'nunchaku-unet'
+ if shared.opts.lora_force_diffusers:
+ return 'diffusers', 'opt-in'
+ if shared.sd_model.__class__.__name__ in force_classes_diffusers:
+ return 'diffusers', 'class-forced'
+ if shared.sd_model_type not in allow_native:
+ return 'diffusers', 'arch-unsupported'
+ if len(shorthash) > 4 and any(x.startswith(shorthash) for x in force_hashes_diffusers):
+ return 'diffusers', 'hash-forced'
+ return 'native', 'default'
def disable_fuse():
diff --git a/modules/lora/network_boft.py b/modules/lora/network_boft.py
new file mode 100644
index 000000000..eccdb30d6
--- /dev/null
+++ b/modules/lora/network_boft.py
@@ -0,0 +1,119 @@
+"""BOFT (Butterfly-OFT) — cascade of butterfly orthogonal factors.
+
+Saves with the same ``oft_blocks`` key as OFT but as a 4-D tensor
+``(boft_m, block_num, block_size, block_size)``. The caller in
+:func:`pipelines.flux.flux2_lora.try_load_oft` discriminates BOFT from
+OFT by ``oft_blocks.ndim == 4``. Math ported from
+``KohakuBlueleaf/LyCORIS/lycoris/modules/boft.py``.
+"""
+
+import torch
+import modules.lora.network as network
+
+
+class ModuleTypeBOFT(network.ModuleType):
+ def create_module(self, net: network.Network, weights: network.NetworkWeights):
+ ob = weights.w.get("oft_blocks")
+ if ob is not None and ob.ndim == 4:
+ return NetworkModuleBOFT(net, weights)
+ return None
+
+
+class NetworkModuleBOFT(network.NetworkModule): # pylint: disable=abstract-method
+ """Butterfly-OFT module: cascade of orthogonal factors.
+
+ Constructor signature mirrors :class:`NetworkModuleOFT` so it slots into
+ the same ``finalize_updown`` pipeline. The ``boft_m``/``block_num``/
+ ``block_size`` triple is read from the saved tensor's shape rather than
+ re-derived via :func:`butterfly_factor`, which keeps loading deterministic
+ even if the upstream factorization heuristic changes.
+ """
+
+ def __init__(self, net: network.Network, weights: network.NetworkWeights):
+ super().__init__(net, weights)
+ self.org_module: list[torch.nn.Module] = [self.sd_module]
+ self.scale = 1.0
+
+ # 4-D oft_blocks: (boft_m, block_num, block_size, block_size)
+ self.oft_blocks = weights.w["oft_blocks"]
+ self.alpha = weights.w["alpha"]
+ self.rescale = weights.w.get("rescale")
+ self.boft_m = self.oft_blocks.shape[0]
+ self.block_num = self.oft_blocks.shape[1]
+ self.block_size = self.oft_blocks.shape[2]
+ self.boft_b = self.block_size
+
+ # Resolve out_dim from the host module — matches NetworkModuleOFT's
+ # discrimination so Linear/Conv2d hosts both work.
+ is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
+ is_conv = type(self.sd_module) in [torch.nn.Conv2d]
+ if is_linear:
+ self.out_dim = self.sd_module.out_features
+ elif is_conv:
+ self.out_dim = self.sd_module.out_channels
+ else:
+ self.out_dim = self.block_num * self.block_size
+
+ # constraint scales with out_dim per LyCORIS BOFT init
+ self.constraint = float(self.alpha) * self.out_dim if self.alpha is not None else 0.0
+
+ def _get_r(self, target: torch.Tensor):
+ """Compute the per-stage Cayley rotations.
+
+ Returns a tensor of shape ``(boft_m, block_num, block_size, block_size)``
+ where each ``r[i]`` is a stack of ``block_num`` orthogonal matrices
+ derived from the i-th butterfly factor via Cayley's parameterization
+ of SO(n): ``R = (I + Q)(I - Q)^-1`` for skew-symmetric ``Q``.
+ """
+ eye = torch.eye(self.block_size, device=target.device, dtype=target.dtype)
+ oft_blocks = self.oft_blocks.to(target.device, dtype=target.dtype)
+ q = oft_blocks - oft_blocks.transpose(-1, -2)
+ if self.constraint > 0:
+ q_norm = torch.norm(q) + 1e-8
+ if q_norm > self.constraint:
+ q = q * self.constraint / q_norm
+ # Inverse needs fp32 to be numerically well-behaved across all dtypes;
+ # cast back to target dtype after.
+ r = (eye + q) @ (eye - q).float().inverse().to(target.dtype)
+ return r
+
+ def _make_weight(self, target: torch.Tensor):
+ """Apply the butterfly cascade to ``target`` and return the transformed weight.
+
+ Direct port of :meth:`ButterflyOFTModule.make_weight` (LyCORIS
+ boft.py:158-191) for the merge-mode (no-bypass) path. ``target`` is the
+ host weight; iteratively reshape to expose the per-stage block layout,
+ einsum-multiply by the stage rotation, then reshape back. The reshape
+ recipe at each stage is what makes the rotations interleave across
+ butterfly partitions, giving the algorithm its O(d log d) parameter
+ density.
+ """
+ m = self.boft_m
+ b = self.boft_b
+ r_b = b // 2
+ r = self._get_r(target)
+ inp = target
+
+ for i in range(m):
+ bi = r[i]
+ g = 2
+ k = 2 ** i * r_b
+ inp = (
+ inp.unflatten(0, (-1, g, k))
+ .transpose(1, 2)
+ .flatten(0, 2)
+ .unflatten(0, (-1, b))
+ )
+ inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
+ inp = (
+ inp.flatten(0, 1).unflatten(0, (-1, k, g)).transpose(1, 2).flatten(0, 2)
+ )
+
+ if self.rescale is not None:
+ inp = inp * self.rescale.to(target.device, dtype=target.dtype)
+ return inp
+
+ def calc_updown(self, target: torch.Tensor):
+ merged = self._make_weight(target)
+ updown = merged - target
+ return self.finalize_updown(updown, target, target.shape)
diff --git a/modules/lora/network_glora.py b/modules/lora/network_glora.py
index ffcb25986..3f00327c3 100644
--- a/modules/lora/network_glora.py
+++ b/modules/lora/network_glora.py
@@ -19,6 +19,7 @@ class NetworkModuleGLora(network.NetworkModule): # pylint: disable=abstract-meth
self.w1b = weights.w["b1.weight"]
self.w2a = weights.w["a2.weight"]
self.w2b = weights.w["b2.weight"]
+ self.dim = self.w1b.shape[0]
def calc_updown(self, target): # pylint: disable=arguments-differ
w1a = self.w1a.to(target.device, dtype=target.dtype)
diff --git a/modules/lora/network_hada.py b/modules/lora/network_hada.py
index 6fc142b3b..1b71b6be6 100644
--- a/modules/lora/network_hada.py
+++ b/modules/lora/network_hada.py
@@ -1,3 +1,4 @@
+import torch
import modules.lora.lyco_helpers as lyco_helpers
import modules.lora.network as network
@@ -44,3 +45,38 @@ class NetworkModuleHada(network.NetworkModule): # pylint: disable=abstract-metho
updown2 = lyco_helpers.rebuild_conventional(w2a, w2b, output_shape)
updown = updown1 * updown2
return self.finalize_updown(updown, target, output_shape)
+
+
+class NetworkModuleHadaChunk(NetworkModuleHada):
+ """LoHA module that returns one row chunk of the Hadamard product.
+
+ Used when a LoHA adapter targets a fused weight (e.g., img_attn.qkv) but the
+ diffusers model exposes separate Q/K/V modules. Slices the row-side of each
+ Hadamard arm (w1a, w2a) at the assigned chunk's row range and computes the
+ partial product. Memory and compute scale linearly with chunk size; no full
+ Hadamard temporary is materialized.
+
+ Tucker (CP-decomposed) LoHAs are not handled here. LyCORIS only saves
+ hada_t1 / hada_t2 for non-1x1 Conv layers, and fused QKV is always Linear,
+ so this combination cannot arise from a conformant trainer.
+ """
+
+ def __init__(self, net, weights, chunk_index, num_chunks):
+ super().__init__(net, weights)
+ self.chunk_index = chunk_index
+ self.num_chunks = num_chunks
+
+ def calc_updown(self, target):
+ w1a = self.w1a.to(target.device, dtype=target.dtype)
+ w1b = self.w1b.to(target.device, dtype=target.dtype)
+ w2a = self.w2a.to(target.device, dtype=target.dtype)
+ w2b = self.w2b.to(target.device, dtype=target.dtype)
+ w1a_chunk = torch.chunk(w1a, self.num_chunks, dim=0)[self.chunk_index].contiguous()
+ w2a_chunk = torch.chunk(w2a, self.num_chunks, dim=0)[self.chunk_index].contiguous()
+ output_shape = [w1a_chunk.size(0), w1b.size(1)]
+ if len(w1b.shape) == 4:
+ output_shape += w1b.shape[2:]
+ updown1 = lyco_helpers.rebuild_conventional(w1a_chunk, w1b, output_shape)
+ updown2 = lyco_helpers.rebuild_conventional(w2a_chunk, w2b, output_shape)
+ updown = updown1 * updown2
+ return self.finalize_updown(updown, target, output_shape)
diff --git a/modules/lora/network_oft.py b/modules/lora/network_oft.py
index e2e61ad45..533cb97b7 100644
--- a/modules/lora/network_oft.py
+++ b/modules/lora/network_oft.py
@@ -55,7 +55,7 @@ class NetworkModuleOFT(network.NetworkModule): # pylint: disable=abstract-method
def calc_updown(self, target):
oft_blocks = self.oft_blocks.to(target.device, dtype=target.dtype)
eye = torch.eye(self.block_size, device=target.device)
- constraint = self.constraint.to(target.device)
+ constraint = self.constraint.to(target.device) if self.constraint is not None else None
if self.is_kohya:
block_Q = oft_blocks - oft_blocks.transpose(1, 2) # ensure skew-symmetric orthogonal matrix
diff --git a/modules/lora/networks.py b/modules/lora/networks.py
index ad95c887e..fe84f3b8a 100644
--- a/modules/lora/networks.py
+++ b/modules/lora/networks.py
@@ -9,6 +9,7 @@ from modules.logger import log, console
applied_layers: list[str] = []
+native_active: bool = False
default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter']
@@ -74,6 +75,8 @@ def network_activate(include=None, exclude=None):
if task is not None and len(applied_layers) == 0:
pbar.remove_task(task) # hide progress bar for no action
+ global native_active # pylint: disable=global-statement
+ native_active = len(l.loaded_networks) > 0
l.timer.activate += time.time() - t0
if l.debug and len(l.loaded_networks) > 0:
log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}')
diff --git a/modules/ltx/ltx_process.py b/modules/ltx/ltx_process.py
index 28c58b2dd..4c5bdc26c 100644
--- a/modules/ltx/ltx_process.py
+++ b/modules/ltx/ltx_process.py
@@ -62,10 +62,12 @@ def _canonical_stage2_kwargs() -> dict:
}
-def _latent_pass(caps, prompt, negative, width, height, frames, steps, guidance_scale, mp4_fps, conditions, image_cond_noise_scale, seed, image=None):
+def _latent_pass(caps, prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask, width, height, frames, steps, guidance_scale, mp4_fps, conditions, image_cond_noise_scale, seed, image=None):
base_args = {
- 'prompt': prompt,
- 'negative_prompt': negative,
+ 'prompt_embeds': prompt_embeds,
+ 'prompt_attention_mask': prompt_attention_mask,
+ 'negative_prompt_embeds': negative_prompt_embeds,
+ 'negative_prompt_attention_mask': negative_prompt_attention_mask,
'width': get_bucket(width),
'height': get_bucket(height),
'num_frames': get_frames(frames),
@@ -93,13 +95,8 @@ def _latent_pass(caps, prompt, negative, width, height, frames, steps, guidance_
base_args['use_cross_timestep'] = True
log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=latent_pass args_keys={list(base_args.keys())}')
result = shared.sd_model(**base_args)
- # video latents strip the batch dim; audio latents keep it so LTX2Pipeline.prepare_audio_latents
- # can rewrap them when re-entered as ndim==4 at Stage 2.
latents = result.frames[0] if hasattr(result, 'frames') else None
- audio_latents = None
- if hasattr(result, 'audio') and result.audio is not None:
- audio_latents = result.audio
- return latents, audio_latents
+ return latents
def run_ltx(task_id,
@@ -269,6 +266,7 @@ def run_ltx(task_id,
vae_type='Default',
vae_tile_frames=16,
)
+ processing.fix_seed(p)
p.scripts = None
p.script_args = None
p.do_not_save_grid = True
@@ -318,7 +316,6 @@ def run_ltx(task_id,
yield None, 'LTX: Generate in progress...'
audio = None
- stage1_audio_latents = None
pixels = None
frames_out = None
needs_latent_path = upsample_enable or refine_enable
@@ -327,10 +324,29 @@ def run_ltx(task_id,
if needs_latent_path:
prompt_final, negative_final, networks = get_prompts(prompt, negative, styles)
extra_networks.activate(p, networks)
- latents, stage1_audio_latents = _latent_pass(
+ # Encode once and reuse across stages; encode_prompt short-circuits when
+ # embeds are passed to __call__. CPU park keeps them off GPU between stages.
+ with devices.inference_context():
+ prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = shared.sd_model.encode_prompt(
+ prompt=prompt_final,
+ negative_prompt=negative_final,
+ do_classifier_free_guidance=True,
+ device=devices.device,
+ )
+ prompt_embeds = prompt_embeds.cpu()
+ prompt_attention_mask = prompt_attention_mask.cpu() if prompt_attention_mask is not None else None
+ negative_prompt_embeds = negative_prompt_embeds.cpu() if negative_prompt_embeds is not None else None
+ negative_prompt_attention_mask = negative_prompt_attention_mask.cpu() if negative_prompt_attention_mask is not None else None
+ # encode_prompt outside pipe.__call__ bypasses the post-forward offload hook;
+ # re-anchor so the text encoder doesn't stay pinned through Stage 1 forward.
+ shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, force=True, silent=True)
+ devices.torch_gc(force=True, reason='ltx:encode')
+ latents = _latent_pass(
caps=caps,
- prompt=prompt_final,
- negative=negative_final,
+ prompt_embeds=prompt_embeds,
+ prompt_attention_mask=prompt_attention_mask,
+ negative_prompt_embeds=negative_prompt_embeds,
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
width=base_w,
height=base_h,
frames=frames,
@@ -339,7 +355,7 @@ def run_ltx(task_id,
mp4_fps=mp4_fps,
conditions=conditions,
image_cond_noise_scale=image_cond_noise_scale if caps.supports_image_cond_noise_scale else None,
- seed=int(seed) if seed is not None else -1,
+ seed=p.seed,
image=p.task_args.get('image'),
)
else:
@@ -348,8 +364,11 @@ def run_ltx(task_id,
yield from abort('Video: process_images returned no frames', ok=True, p=p)
return
pixels = processed.images
- if getattr(processed, 'audio', None) is not None:
- audio = processed.audio
+ raw_audio = getattr(processed, 'audio', None)
+ if raw_audio is not None:
+ # Strip batch dim from (B, 2, N); write_audio expects (2, N) for the
+ # transpose-to-interleaved path used by AAC s16.
+ audio = raw_audio[0].float().cpu() if raw_audio.ndim == 3 else raw_audio.float().cpu()
latents = None
except AssertionError as e:
yield from abort(e, ok=True, p=p)
@@ -388,7 +407,7 @@ def run_ltx(task_id,
up_args = {
'width': final_w,
'height': final_h,
- 'generator': get_generator(int(seed) if seed is not None else -1),
+ 'generator': get_generator(p.seed),
'output_type': 'latent',
}
if latents.ndim == 4:
@@ -409,7 +428,7 @@ def run_ltx(task_id,
'height': final_h,
'num_frames': get_frames(frames),
'latents_normalized': False,
- 'generator': get_generator(int(seed) if seed is not None else -1),
+ 'generator': get_generator(p.seed),
'output_type': 'latent',
}
if latents.ndim == 4:
@@ -436,13 +455,15 @@ def run_ltx(task_id,
# Refine is terminal: let the pipe decode internally so the final VAE pass runs inside
# the same offload/cudnn context as a normal generation (matches Generic Video tab).
refine_args = {
- 'prompt': prompt_final,
- 'negative_prompt': negative_final,
+ 'prompt_embeds': prompt_embeds,
+ 'prompt_attention_mask': prompt_attention_mask,
+ 'negative_prompt_embeds': negative_prompt_embeds,
+ 'negative_prompt_attention_mask': negative_prompt_attention_mask,
'width': final_w,
'height': final_h,
'num_frames': get_frames(frames),
'num_inference_steps': steps,
- 'generator': get_generator(int(seed) if seed is not None else -1),
+ 'generator': get_generator(p.seed),
'callback_on_step_end': diffusers_callback,
'output_type': 'pil',
}
@@ -457,14 +478,14 @@ def run_ltx(task_id,
# Thread Stage-1 I2V init image through Stage 2 so first-frame identity survives refine.
if caps.is_i2v and caps.repo_cls_name in ('LTXImageToVideoPipeline', 'LTX2ImageToVideoPipeline') and p.task_args.get('image') is not None:
refine_args['image'] = p.task_args['image']
- # Thread Stage-1 audio latents into Stage 2 on 2.x. The video branch cross-attends
- # audio every layer; letting prepare_audio_latents fall back to fresh noise biases
- # the video branch off-distribution (desaturated output on distilled 2.x).
- if caps.family == '2.x':
- if stage1_audio_latents is not None:
- refine_args['audio_latents'] = stage1_audio_latents.to(device=devices.device)
- if caps.use_cross_timestep:
- refine_args['use_cross_timestep'] = True
+ if caps.family == '2.x' and caps.use_cross_timestep:
+ refine_args['use_cross_timestep'] = True
+ # output_type='latent' skips the post-loop audio_vae + vocoder pass when audio
+ # is unwanted; per-step audio cross-attention still runs for video conditioning.
+ # Internal video decode is also skipped; vae_decode below picks it up.
+ want_audio = caps.supports_audio and audio_enable
+ if not want_audio:
+ refine_args['output_type'] = 'latent'
saved_scheduler_stage2 = None
try:
@@ -502,10 +523,14 @@ def run_ltx(task_id,
yield None, 'LTX: Refine in progress...'
try:
result = shared.sd_model(latents=latents, **refine_args)
- pixels = result.frames[0] if hasattr(result, 'frames') else None
- if hasattr(result, 'audio') and result.audio is not None:
- audio = result.audio[0].float().cpu()
- latents = None
+ out = result.frames[0] if hasattr(result, 'frames') else None
+ if want_audio:
+ pixels = out
+ if hasattr(result, 'audio') and result.audio is not None:
+ audio = result.audio[0].float().cpu()
+ latents = None
+ else:
+ latents = out
except AssertionError as e:
yield from abort(e, ok=True, p=p)
return
@@ -533,14 +558,15 @@ def run_ltx(task_id,
extra_networks.deactivate(p)
if needs_latent_path and latents is not None:
- # Only reached on upsample-without-refine; refine decodes through the pipe and nulls latents.
+ # Decode any path that leaves latents intact: upsample-without-refine, or
+ # refine with output_type='latent' (audio_enable=False).
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'], force=True, silent=True)
devices.torch_gc(force=True, reason='ltx:vae')
yield None, 'LTX: VAE decode in progress...'
try:
if torch.is_tensor(latents):
# 0.9.x returns raw latents with output_type='latent'; 2.x pre-denormalizes.
- frames_out = vae_decode(latents, decode_timestep if caps.supports_decode_timestep else 0.0, int(seed) if seed is not None else -1, denormalize=caps.family == '0.9')
+ frames_out = vae_decode(latents, decode_timestep if caps.supports_decode_timestep else 0.0, p.seed, denormalize=caps.family == '0.9')
else:
frames_out = latents
except AssertionError as e:
diff --git a/modules/ltx/ltx_ui.py b/modules/ltx/ltx_ui.py
index 896a89cf7..7385d5351 100644
--- a/modules/ltx/ltx_ui.py
+++ b/modules/ltx/ltx_ui.py
@@ -109,7 +109,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
audio_accordion = gr.Accordion(open=False, label="Audio", elem_id='ltx_audio_accordion', visible=False)
with audio_accordion:
with gr.Row():
- audio_enable = gr.Checkbox(label='LTX enable audio', value=False, elem_id="ltx_audio_enable")
+ audio_enable = gr.Checkbox(label='LTX save audio', value=True, elem_id="ltx_audio_enable")
with gr.Column(elem_id='ltx-output-column', scale=2) as _column_output:
with gr.Row():
diff --git a/modules/masking.py b/modules/masking.py
index 9b5b193fa..0accc7eef 100644
--- a/modules/masking.py
+++ b/modules/masking.py
@@ -290,9 +290,9 @@ def run_rembg(input_image: Image.Image, input_mask: np.ndarray):
'alpha_matting_foreground_threshold': 240,
'alpha_matting_background_threshold': 10,
'alpha_matting_erode_size': int(opts.mask_erode * 40),
- 'session': rembg.new_session(opts.model),
+ 'session': rembg.new_session(opts.model), # pylint: disable=c-extension-no-member
}
- mask = rembg.remove(**args)
+ mask = rembg.remove(**args) # pylint: disable=c-extension-no-member
mask = np.array(mask)
if input_mask is None:
input_mask = np.zeros(mask.shape, dtype='uint8')
@@ -385,7 +385,9 @@ def outpaint(input_image: Image.Image, outpaint_type: str = 'Edge'):
def run_mask(input_image: Image.Image, input_mask: Image.Image | None = None, return_type: str | None = None, mask_blur: int | None = None, mask_padding: int | None = None, invert=None):
- if isinstance(input_image, list) and len(input_image) > 0:
+ if input_image is None:
+ return input_mask
+ elif isinstance(input_image, list) and len(input_image) > 0:
input_image = input_image[0]
elif isinstance(input_image, dict):
input_mask = input_image.get('mask', None)
@@ -401,7 +403,10 @@ def run_mask(input_image: Image.Image, input_mask: Image.Image | None = None, re
debug(f'Run mask: fn={fn}') # pylint: disable=protected-access
debug(f'Run mask: opts={opts}') # pylint: disable=protected-access
- size = min(input_image.width, input_image.height)
+ try:
+ size = min(input_image.width, input_image.height)
+ except Exception:
+ return input_mask
if invert is not None:
opts.invert = invert
@@ -574,13 +579,13 @@ def create_segment_ui():
return controls
-def bind_controls(image_controls: list[gr.Image], preview_image: gr.Image, output_image: gr.Image):
+def bind_controls(image_controls: list[gr.Image], output_image: gr.Image):
for image_control in image_controls:
- btn_mask.click(run_mask, inputs=[image_control], outputs=[preview_image])
+ btn_mask.click(run_mask, inputs=[image_control], outputs=[output_image])
btn_lama.click(run_lama, inputs=[image_control], outputs=[output_image])
- image_control.edit(fn=run_mask_live, inputs=[image_control], outputs=[preview_image])
+ image_control.edit(fn=run_mask_live, inputs=[image_control], outputs=[output_image])
for control in controls:
- control.change(fn=run_mask_live, inputs=[image_control], outputs=[preview_image])
+ control.change(fn=run_mask_live, inputs=[image_control], outputs=[output_image])
def process_kanvas(kanvas_data):
diff --git a/modules/modeldata.py b/modules/modeldata.py
index 536c2380e..4662f8165 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -48,6 +48,8 @@ def get_model_type(pipe):
model_type = 'chroma'
elif "Flux2" in name:
model_type = 'f2'
+ elif "UltraFlux" in name:
+ model_type = 'ultraflux'
elif "Flux" in name or "Flex1" in name or "Flex2" in name:
model_type = 'f1'
elif "ZImage" in name or "Z-Image" in name:
@@ -70,9 +72,11 @@ def get_model_type(pipe):
model_type = 'sana'
elif 'VIBE' in name:
model_type = 'sana'
+ elif "HiDreamO1" in name:
+ model_type = 'o1'
elif "HiDream" in name:
model_type = 'h1'
- elif "AnimaTextToImage" in name:
+ elif name.startswith("Anima") and "AnimateDiff" not in name:
model_type = 'anima'
elif "Cosmos2TextToImage" in name:
model_type = 'cosmos'
@@ -108,6 +112,8 @@ def get_model_type(pipe):
model_type = 'glmimage'
elif 'Step1XEdit' in name:
model_type = 'step1x_edit'
+ elif 'JoyImageEdit' in name:
+ model_type = 'joy'
elif 'OvisImage' in name or 'Ovis-Image' in name:
model_type = 'ovis'
elif 'Wan' in name:
diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py
index 82c9740a8..87d6a746a 100644
--- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py
+++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py
@@ -29,7 +29,7 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip
feature_extractor: Any,
requires_safety_checker: bool = True
):
- super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker)
+ super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker) # pylint: disable=too-many-function-args
self.image_processor = VaeImageProcessor(vae_scale_factor=64)
def __call__(
@@ -72,7 +72,7 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip
image = self.image_processor.preprocess(image).cpu().numpy()
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
diff --git a/modules/options_handler.py b/modules/options_handler.py
index 685363e46..ffc3080ea 100644
--- a/modules/options_handler.py
+++ b/modules/options_handler.py
@@ -91,7 +91,7 @@ class Options:
return
return super().__setattr__(key, value) # pylint: disable=super-with-arguments
- def set(self, key, value):
+ def set(self, key, value, force=False):
"""sets an option and calls its onchange callback, returning True if the option changed and False otherwise"""
if key in self.secrets:
oldval = self.secrets.get(key, None)
@@ -103,7 +103,7 @@ class Options:
else:
log.warning(f'Settings: key={key} value={value} unknown')
return False
- if oldval == value:
+ if oldval == value and not force:
return False
try:
setattr(self, key, value)
diff --git a/modules/pag/pipe_sd.py b/modules/pag/pipe_sd.py
index 4393a24d3..5f28400cf 100644
--- a/modules/pag/pipe_sd.py
+++ b/modules/pag/pipe_sd.py
@@ -1057,11 +1057,11 @@ class StableDiffusionPAGPipeline(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: List[int] | None = None,
guidance_scale: float = 7.5,
pag_scale: float = 0.0,
pag_adaptive_scaling: float = 0.0,
@@ -1321,10 +1321,10 @@ class StableDiffusionPAGPipeline(
up_layers[int(drop_layer[1])].processor = replace_processor
else:
raise ValueError(f"Invalid layer type: {drop_layer[0]}")
- except IndexError:
+ except IndexError as e:
raise ValueError(
f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers."
- )
+ ) from e
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
@@ -1452,10 +1452,10 @@ class StableDiffusionPAGPipeline(
up_layers[int(drop_layer[1])].processor = AttnProcessor2_0()
else:
raise ValueError(f"Invalid layer type: {drop_layer[0]}")
- except IndexError:
+ except IndexError as e:
raise ValueError(
f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers."
- )
+ ) from e
if not return_dict:
return (image, has_nsfw_concept)
diff --git a/modules/pag/pipe_sdxl.py b/modules/pag/pipe_sdxl.py
index 653fecf30..02350cd0a 100644
--- a/modules/pag/pipe_sdxl.py
+++ b/modules/pag/pipe_sdxl.py
@@ -1079,19 +1079,19 @@ class StableDiffusionXLPAGPipeline(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: List[int] | None = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
pag_scale: float = 0.0,
pag_adaptive_scaling: float = 0.0,
pag_drop_rate: float = 0.5,
pag_applied_layers: List[str] = ['mid'], #['down', 'mid', 'up']
- pag_applied_layers_index: List[str] = None, #['d4', 'd5', 'm0']
+ pag_applied_layers_index: List[str] | None = None, #['d4', 'd5', 'm0']
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
@@ -1507,10 +1507,10 @@ class StableDiffusionXLPAGPipeline(
up_layers[layer_number].processor = replace_processor
else:
raise ValueError(f"Invalid layer type: {drop_layer[0]}")
- except IndexError:
+ except IndexError as e:
raise ValueError(
f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers."
- )
+ ) from e
elif self.pag_applied_layers:
drop_full_layers = self.pag_applied_layers
for drop_full_layer in drop_full_layers:
@@ -1526,10 +1526,10 @@ class StableDiffusionXLPAGPipeline(
up_layer.processor = replace_processor
else:
raise ValueError(f"Invalid layer type: {drop_full_layer}")
- except IndexError:
+ except IndexError as e:
raise ValueError(
f"Invalid layer index: {drop_full_layer}. Available layers are: down, mid and up. If you need to specify each layer index, you can use `pag_applied_layers_index`"
- )
+ ) from e
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
@@ -1678,8 +1678,8 @@ class StableDiffusionXLPAGPipeline(
up_layers[layer_number].processor = AttnProcessor2_0()
else:
raise ValueError(f"Invalid layer type: {drop_layer[0]}")
- except IndexError:
- raise ValueError(f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers.")
+ except IndexError as e:
+ raise ValueError(f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers.") from e
elif self.pag_applied_layers:
drop_full_layers = self.pag_applied_layers
for drop_full_layer in drop_full_layers:
@@ -1695,6 +1695,6 @@ class StableDiffusionXLPAGPipeline(
up_layer.processor = AttnProcessor2_0()
else:
raise ValueError(f"Invalid layer type: {drop_full_layer}")
- except IndexError:
- raise ValueError(f"Invalid layer index: {drop_full_layer}. Available layers are: down, mid and up. If you need to specify each layer index, you can use `pag_applied_layers_index`")
+ except IndexError as e:
+ raise ValueError(f"Invalid layer index: {drop_full_layer}. Available layers are: down, mid and up. If you need to specify each layer index, you can use `pag_applied_layers_index`") from e
return StableDiffusionXLPipelineOutput(images=image)
diff --git a/modules/postprocess/aurasr_arch.py b/modules/postprocess/aurasr_arch.py
index 2e5eb0ea6..031996a8b 100644
--- a/modules/postprocess/aurasr_arch.py
+++ b/modules/postprocess/aurasr_arch.py
@@ -14,6 +14,7 @@ import torch.nn.functional as F
from einops import rearrange, repeat, reduce
from einops.layers.torch import Rearrange
+import itertools
def get_same_padding(size, kernel, dilation, stride):
@@ -262,7 +263,7 @@ class LinearAttention(nn.Module):
self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), RMSNorm(dim))
def forward(self, x):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
x = self.norm(x)
@@ -296,7 +297,7 @@ class Attention(nn.Module):
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
def forward(self, x):
- b, c, h, w = x.shape
+ _b, _c, h, w = x.shape
x = self.norm(x)
qkv = self.to_qkv(x).chunk(3, dim=1)
@@ -487,8 +488,8 @@ class UnetUpsampler(torch.nn.Module):
down_dims = [init_down_dim, *map(lambda m: dim * m, down_dim_mults)]
self.init_conv = nn.Conv2d(input_channels, init_down_dim, 7, padding=3)
- up_in_out = list(zip(up_dims[:-1], up_dims[1:]))
- down_in_out = list(zip(down_dims[:-1], down_dims[1:]))
+ up_in_out = list(itertools.pairwise(up_dims))
+ down_in_out = list(itertools.pairwise(down_dims))
block_klass = partial(
ResnetBlock,
@@ -507,7 +508,7 @@ class UnetUpsampler(torch.nn.Module):
block_count = 6
- for ind, (
+ for _, (
(dim_in, dim_out),
layer_full_attn,
layer_attn_depth,
@@ -515,7 +516,7 @@ class UnetUpsampler(torch.nn.Module):
attn_klass = FullAttention if layer_full_attn else LinearTransformer
blocks = []
- for i in range(block_count):
+ for _i in range(block_count):
blocks.append(block_klass(dim_in, dim_in))
self.downs.append(
@@ -569,7 +570,7 @@ class UnetUpsampler(torch.nn.Module):
blocks = []
input_dim = dim_in * 2 if ind < len(down_in_out) else dim_in
- for i in range(block_count):
+ for _i in range(block_count):
blocks.append(block_klass(input_dim, dim_in))
self.ups.append(
@@ -708,7 +709,7 @@ class UnetUpsampler(torch.nn.Module):
def tile_image(image, chunk_size=64):
- c, h, w = image.shape
+ _c, h, w = image.shape
h_chunks = ceil(h / chunk_size)
w_chunks = ceil(w / chunk_size)
tiles = []
@@ -787,12 +788,12 @@ class AuraSR:
try:
from safetensors.torch import load_file
checkpoint = load_file(hf_model_path / "model.safetensors" if not Path(model_id).is_file() else model_id)
- except ImportError:
+ except ImportError as e:
raise ImportError(
"The safetensors library is not installed. "
"Please install it with `pip install safetensors` "
"or use `use_safetensors=False` to load the model with PyTorch."
- )
+ ) from e
else:
checkpoint = torch.load(hf_model_path / "model.ckpt" if not Path(model_id).is_file() else model_id)
diff --git a/modules/postprocess/swinir_model_arch.py b/modules/postprocess/swinir_model_arch.py
index 4b306433d..8064e0fef 100644
--- a/modules/postprocess/swinir_model_arch.py
+++ b/modules/postprocess/swinir_model_arch.py
@@ -238,7 +238,7 @@ class SwinTransformerBlock(nn.Module):
def forward(self, x, x_size):
H, W = x_size
- B, L, C = x.shape
+ B, _L, C = x.shape
# assert L == H * W, "input feature has wrong size"
shortcut = x
@@ -559,7 +559,7 @@ class PatchUnEmbed(nn.Module):
self.embed_dim = embed_dim
def forward(self, x, x_size):
- B, HW, C = x.shape
+ B, _HW, _C = x.shape
x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C
return x
diff --git a/modules/postprocess/swinir_model_arch_v2.py b/modules/postprocess/swinir_model_arch_v2.py
index d61e92668..00577c8ba 100644
--- a/modules/postprocess/swinir_model_arch_v2.py
+++ b/modules/postprocess/swinir_model_arch_v2.py
@@ -266,7 +266,7 @@ class SwinTransformerBlock(nn.Module):
def forward(self, x, x_size):
H, W = x_size
- B, L, C = x.shape
+ B, _L, C = x.shape
#assert L == H * W, "input feature has wrong size"
shortcut = x
@@ -476,7 +476,7 @@ class PatchEmbed(nn.Module):
self.norm = None
def forward(self, x):
- B, C, H, W = x.shape
+ _B, _C, _H, _W = x.shape
# FIXME look at relaxing size constraints
# assert H == self.img_size[0] and W == self.img_size[1],
# f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
@@ -591,7 +591,7 @@ class PatchUnEmbed(nn.Module):
self.embed_dim = embed_dim
def forward(self, x, x_size):
- B, HW, C = x.shape
+ B, _HW, _C = x.shape
x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C
return x
diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py
index e70df763f..b89ec5aac 100644
--- a/modules/postprocess/yolo.py
+++ b/modules/postprocess/yolo.py
@@ -178,7 +178,7 @@ class YoloRestorer(Detailer):
max_size = opt_max if 0 < opt_max <= 1 else 1
if x_size >= min_size and y_size >=min_size and x_size <= max_size and y_size <= max_size:
if mask:
- if detailer_opt(p, 'detailer_segmentation', 'detailer_seg') and seg is not None:
+ if detailer_opt(p, 'detailer_segmentation') and seg is not None:
masked = seg
else:
masked = Image.new('L', image.size, 0)
@@ -267,7 +267,7 @@ class YoloRestorer(Detailer):
color = (0, 190, 190)
log.debug(f'Detailer: draw={items}')
for i, item in enumerate(items):
- if detailer_opt(p, 'detailer_segmentation', 'detailer_seg') and item.mask is not None:
+ if detailer_opt(p, 'detailer_segmentation') and item.mask is not None:
mask = item.mask.convert('L')
else:
mask = Image.new('L', image.size, 0)
@@ -427,9 +427,12 @@ class YoloRestorer(Detailer):
pc.negative_prompt = negative_lines[i*len(items)+j]
pc.prompts = [pc.prompt]
pc.negative_prompts = [pc.negative_prompt]
- pc.prompts, pc.network_data = extra_networks.parse_prompts(pc.prompts)
- extra_networks.activate(pc, pc.network_data)
- log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={detailer_opt(p, "detailer_segmentation", "detailer_seg")} prompt="{pc.prompt}"')
+ pc.prompts, pc.network_data = extra_networks.parse_prompts(pc.prompts, pc.network_data)
+ pc.disable_extra_networks = True # disable processing_diffusers from handling network activation since its handled here
+ network_same = len(p.network_data.values()) == len(pc.network_data.values()) and all(x == y for x, y in zip(p.network_data.values(), pc.network_data.values()))
+ if not network_same:
+ extra_networks.activate(pc, pc.network_data)
+ log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={detailer_opt(p, "detailer_segmentation")} network={network_same} prompt="{pc.prompt}"')
pc.init_images = [image]
pc.image_mask = [item.mask]
pc.overlay_images = []
@@ -443,7 +446,8 @@ class YoloRestorer(Detailer):
# process
jobid = shared.state.begin('Detailer')
pp = processing.process_images_inner(pc)
- extra_networks.deactivate(pc, force=True)
+ if not network_same:
+ extra_networks.deactivate(pc, force=True)
shared.sd_model.fail_on_switch_error = False
shared.state.end(jobid)
diff --git a/modules/processing.py b/modules/processing.py
index 0306b20a8..188ae6551 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -135,7 +135,7 @@ def get_processed(*args, **kwargs):
return processed
-def process_images(p: StableDiffusionProcessing) -> Processed:
+def process_images(p: StableDiffusionProcessing) -> Processed | None:
timer.process.reset()
debug(f'Process images: class={p.__class__.__name__} {vars(p)}')
if shared.sd_model is None:
diff --git a/modules/processing_args.py b/modules/processing_args.py
index 6f3305e62..ed7150f4b 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -128,15 +128,18 @@ def task_specific_kwargs(p, model):
'width': width,
}
+ fake_i2i = ['QwenImageEditPipeline', 'QwenImageEditPlusPipeline', 'WanImageToVideoPipeline', 'ChronoEditPipeline']
+ can_i2i = ['QwenImageEditPipeline', 'QwenImageEditPlusPipeline', 'QwenImageLayeredPipeline', 'Kandinsky5I2IPipeline', 'QwenImageLayeredPipeline', 'WanImageToVideoPipeline','ChronoEditPipeline', 'GoogleNanoBananaPipeline', 'GlmImagePipeline', 'Step1XEditPipeline']
+
# model specific args
- if ('QwenImageEdit' in model_cls) and (p.init_images is None or len(p.init_images) == 0):
- task_args['image'] = [Image.new('RGB', (p.width, p.height), (0, 0, 0))] # monkey-patch so qwen-image-edit pipeline does not error-out on t2i
- if ('QwenImageEditPlusPipeline' in model_cls) and (p.init_control is not None) and (len(p.init_control) > 0):
- task_args['image'] += p.init_control
- if ('QwenImageLayeredPipeline' in model_cls) and (p.init_images is not None) and (len(p.init_images) > 0):
- task_args['image'] = p.init_images[0].convert('RGBA')
- if ('Flux2' in model_cls) and (p.init_control is not None) and (len(p.init_control) > 0):
- task_args['image'] += p.init_control
+ if (model_cls in fake_i2i) and (len(getattr(p, 'init_images', [])) == 0):
+ log.debug(f'Model init: cls={model_cls} image=blank')
+ p.init_images = [Image.new('RGB', (p.width, p.height), (0, 0, 0))] # monkey-patch so i2i pipeline does not error-out on t2i
+ if (model_cls in can_i2i) and (len(getattr(p, 'init_images', [])) > 0):
+ task_args['image'] = p.init_images
+
+ if ('QwenImageLayeredPipeline' in model_cls) and (task_args.get('image', None) is not None):
+ task_args['image'] = [i.convert('RGBA') for i in task_args['image']]
if ('LatentConsistencyModelPipeline' in model_cls) and (len(p.init_images) > 0):
p.ops.append('lcm')
init_latents = [processing_vae.vae_encode(image, model=shared.sd_model, vae_type=p.vae_type).squeeze(dim=0) for image in p.init_images]
@@ -148,17 +151,8 @@ def task_specific_kwargs(p, model):
'width': p.width,
'height': p.height,
}
- if ('WanImageToVideoPipeline' in model_cls) or ('ChronoEditPipeline' in model_cls):
- if (p.init_images is not None) and (len(p.init_images) > 0):
- task_args['image'] = p.init_images[0]
- else:
- task_args['image'] = Image.new('RGB', (p.width, p.height), (0, 0, 0)) # monkey-patch so wan-i2i pipeline does not error-out on t2i
if ('WanVACEPipeline' in model_cls) and (p.init_images is not None) and (len(p.init_images) > 0):
task_args['reference_images'] = p.init_images
- if ('GoogleNanoBananaPipeline' in model_cls) and (p.init_images is not None) and (len(p.init_images) > 0):
- task_args['image'] = p.init_images[0]
- if ('GlmImagePipeline' in model_cls) and (p.init_images is not None) and (len(p.init_images) > 0):
- task_args['image'] = p.init_images
if 'BlipDiffusionPipeline' in model_cls:
if len(p.init_images) == 0:
log.error('BLiP diffusion requires init image')
@@ -259,6 +253,9 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
args['guidance_scale'] = p.cfg_scale
if 'img_guidance_scale' in possible and hasattr(p, 'image_cfg_scale') and p.image_cfg_scale is not None and p.image_cfg_scale > 0:
args['img_guidance_scale'] = p.image_cfg_scale
+ if getattr(getattr(model, 'config', None), 'is_distilled', False) and args.get('guidance_scale', 0) > 1 and not getattr(p, 'distilled_warned', False):
+ log.warning(f'Pipeline: cls={model.__class__.__name__} distilled=True cfg_scale={args["guidance_scale"]} ignored, forced to 1')
+ p.distilled_warned = True
if 'generator' in possible:
generator = get_generator(p)
args['generator'] = generator
@@ -272,7 +269,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
kwargs['output_type'] = 'np' # only set latent if model has vae
# model specific
- if 'Kandinsky' in model.__class__.__name__ or 'Cosmos2' in model.__class__.__name__ or 'Anima' in model.__class__.__name__ or 'OmniGen2' in model.__class__.__name__:
+ if 'Kandinsky' in model.__class__.__name__ or 'Cosmos2' in model.__class__.__name__ or 'OmniGen2' in model.__class__.__name__:
kwargs['output_type'] = 'np' # only set latent if model has vae
if 'StableCascade' in model.__class__.__name__:
kwargs.pop("guidance_scale") # remove
diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py
index e1c2c19ba..c9689c6e5 100644
--- a/modules/processing_callbacks.py
+++ b/modules/processing_callbacks.py
@@ -59,6 +59,23 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
torch.xpu.synchronize(devices.device)
elif devices.backend in {"cuda", "zluda", "rocm"}:
torch.cuda.synchronize(devices.device)
+
+ if shared.state.paused:
+ log.debug('Sampling paused')
+ while shared.state.paused:
+ if shared.state.interrupted or shared.state.skipped:
+ raise AssertionError('Interrupted...')
+ time.sleep(0.1)
+
+ image = kwargs.get('image', None)
+ if image is not None:
+ shared.state.current_image = image
+ shared.state.current_latent = None
+ shared.state.step() # increase step
+ shared.state.preview_job = -1 # indicate that preview image has changed
+ debug_callback(f'Callback: step={step} timestep={timestep} image={image if image is not None else None} kwargs={list(kwargs)}')
+ return kwargs
+
latents = kwargs.get('latents', None)
if debug:
debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}')
@@ -67,12 +84,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
shared.state.step()
if shared.state.interrupted or shared.state.skipped:
raise AssertionError('Interrupted...')
- if shared.state.paused:
- log.debug('Sampling paused')
- while shared.state.paused:
- if shared.state.interrupted or shared.state.skipped:
- raise AssertionError('Interrupted...')
- time.sleep(0.1)
if latents is None:
return kwargs
elif shared.opts.nan_skip:
@@ -139,7 +150,9 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
width = getattr(p, 'width', 1024)
height = getattr(p, 'height', 1024)
latents = kwargs['latents']
- if len(latents.shape) == 3: # packed format [B, seq_len, patch_channels]
+ if len(latents.shape) == 4:
+ latents = pipe._unpatchify_latents(latents) # [B, C*4, h/2, w/2] -> [B, C, h, w] # pylint: disable=protected-access
+ elif len(latents.shape) == 3: # packed format [B, seq_len, patch_channels]
b, seq_len, patch_ch = latents.shape
channels = patch_ch // 4 # 4 = 2x2 patch
h_patches = height // vae_scale // 2
diff --git a/modules/processing_class.py b/modules/processing_class.py
index 9e1ed3e4f..581914204 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -276,8 +276,10 @@ class StableDiffusionProcessing:
tome_ratio: float | None = None,
todo_ratio: float | None = None,
# overrides
+ skip_processing: bool = False,
override_settings_restore_afterwards: bool = True,
override_settings: dict[str, Any] | None = None,
+ network_data: dict | None = None,
# metadata
# extra_generation_params: Dict[Any, Any] = {},
# task_args: Dict[str, Any] = {},
@@ -321,7 +323,7 @@ class StableDiffusionProcessing:
self.negative_prompt_attention_masks = []
self.disable_extra_networks = False
self.iteration = 0
- self.network_data = {}
+ self.network_data = network_data or {}
# initializers
self.prompt = prompt
@@ -491,6 +493,7 @@ class StableDiffusionProcessing:
self.scale_by_before = scale_by_before
self.scale_by_after = scale_by_after
self.scale_by_mask = scale_by_mask
+ self.skip_processing = skip_processing
# special handled items
if firstphase_width != 0 or firstphase_height != 0:
diff --git a/modules/processing_correction.py b/modules/processing_correction.py
index 1c3751bb9..9477960dd 100644
--- a/modules/processing_correction.py
+++ b/modules/processing_correction.py
@@ -212,6 +212,8 @@ def _count_steps_below(pipe, threshold):
def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False, step: int = 0):
+ if pipe and pipe.__class__.__name__ in ['HiDreamO1Pipeline', 'HiDreamO1ImagePipeline']:
+ return kwargs
if initial:
if not any([p.hdr_clamp, p.hdr_mode, p.hdr_maximize, p.hdr_sharpen, p.hdr_color, p.hdr_brightness, p.hdr_tint_ratio]):
p.correction_skip = True
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index b7b1030ef..562d31fe4 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -242,7 +242,7 @@ def process_base(p: processing.StableDiffusionProcessing):
def process_hires(p: processing.StableDiffusionProcessing, output):
# optional second pass
- if (output is None) or (output.images is None):
+ if (output is None) or not hasattr(output, 'images') or (output.images is None):
return output
if p.enable_hr:
jobid = shared.state.begin('Hires')
@@ -368,7 +368,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
def process_refine(p: processing.StableDiffusionProcessing, output):
# optional refiner pass or decode
- if (output is None) or (output.images is None):
+ if (output is None) or not hasattr(output, 'images') or (output.images is None):
return output
if is_refiner_enabled(p):
if shared.opts.samples_save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
@@ -452,17 +452,32 @@ def process_refine(p: processing.StableDiffusionProcessing, output):
return output
+class AudioFrameList(list):
+ """list subclass with an audio attribute. Plain lists can't hold attributes,
+ so process_decode uses this when the pipeline output includes audio."""
+ audio = None
+
+
+def attach_audio(results, audio):
+ if audio is None:
+ return results
+ wrapped = AudioFrameList(results if isinstance(results, list) else list(results))
+ wrapped.audio = audio
+ return wrapped
+
+
def process_decode(p: processing.StableDiffusionProcessing, output):
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'])
if output is not None:
if hasattr(output, 'bytes') and output.bytes is not None:
log.debug(f'Generated: bytes={len(output.bytes)}')
return output
+ audio = getattr(output, 'audio', None)
if not hasattr(output, 'images') and hasattr(output, 'frames'):
log.debug(f'Generated: frames={len(output.frames[0])}')
output.images = output.frames[0]
if output.images is not None and len(output.images) > 0 and isinstance(output.images[0], Image.Image):
- return output.images
+ return attach_audio(output.images, audio)
model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner
if not hasattr(model, 'vae'):
if hasattr(model, 'pipe') and hasattr(model.pipe, 'vae'):
@@ -506,8 +521,9 @@ def process_decode(p: processing.StableDiffusionProcessing, output):
results = []
else:
log.warning('Processing: no results')
+ audio = None
results = []
- return results
+ return attach_audio(results, audio)
def update_pipeline(sd_model, p: processing.StableDiffusionProcessing):
diff --git a/modules/processing_grading.py b/modules/processing_grading.py
index 055e72de7..ec9fada42 100644
--- a/modules/processing_grading.py
+++ b/modules/processing_grading.py
@@ -181,10 +181,16 @@ def _apply_color_temp(img: torch.Tensor, kelvin: float) -> torch.Tensor:
def _apply_lut(image: Image.Image, lut_cube_file: str, strength: float) -> Image.Image:
"""Apply .cube LUT file via pillow-lut-tools."""
- if not lut_cube_file or not os.path.isfile(lut_cube_file):
+ if not lut_cube_file:
return image
- pillow_lut = _ensure_pillow_lut()
+ if hasattr(lut_cube_file, "name"):
+ lut_cube_file = lut_cube_file.name
+ if hasattr(lut_cube_file, "as_posix"):
+ lut_cube_file = lut_cube_file.as_posix()
try:
+ if not os.path.isfile(lut_cube_file):
+ return image
+ pillow_lut = _ensure_pillow_lut()
cube = pillow_lut.load_cube_file(lut_cube_file)
if strength != 1.0:
cube = pillow_lut.amplify_lut(cube, strength)
diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py
index 7e4b1a4c6..55e2ccb42 100644
--- a/modules/processing_helpers.py
+++ b/modules/processing_helpers.py
@@ -562,12 +562,17 @@ def apply_circular(enable: bool, model):
def save_intermediate(p, latents, suffix):
- for i in range(len(latents)):
- from modules.processing import create_infotext
- info=create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i)
+ from modules.processing import create_infotext
+ from modules.image import convert
+ is_latent = torch.is_tensor(latents) and latents.shape[-1] != 3
+ if is_latent:
decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, output_type='pil', vae_type=p.vae_type, width=p.width, height=p.height)
- for j in range(len(decoded)):
- images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
+ else:
+ items = latents if isinstance(latents, list) else ([latents[j] for j in range(latents.shape[0])] if hasattr(latents, 'shape') else [latents])
+ decoded = [convert.to_pil(img) if not hasattr(img, 'width') else img for img in items]
+ for i in range(len(decoded)):
+ info = create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i)
+ images.save_image(decoded[i], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
def update_sampler(p, sd_model, second_pass=False):
diff --git a/modules/processing_prompt.py b/modules/processing_prompt.py
index 212a6cc58..32ad88ef1 100644
--- a/modules/processing_prompt.py
+++ b/modules/processing_prompt.py
@@ -152,7 +152,7 @@ def set_prompt(p,
negative_attention_masks = prompt_parser_diffusers.embedder('negative_prompt_attention_masks')
if negative_embeds is None:
- log.warning('Prompt parser encode: empty negative prompt embeds')
+ # log.warning('Prompt parser encode: empty negative prompt embeds')
prompt_parser_diffusers.embedder = None
args = set_fallback_prompt(args, possible, prompts=None, negative_prompts=negative_prompts, prompts_2=None, negative_prompts_2=None)
prompt_attention = 'fixed'
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index d35a1d15b..b793140b5 100644
--- a/modules/processing_vae.py
+++ b/modules/processing_vae.py
@@ -68,7 +68,7 @@ def full_vqgan_decode(latents, model):
decoded = []
# delete vae after OpenVINO compile
- if 'VAE' in shared.opts.cuda_compile and shared.compiled_model_state.first_pass_vae and (shared.opts.cuda_compile_backend == "openvino_fx" or shared.opts.cuda_compile_backend == "openvino"):
+ if ('VAE' in shared.opts.cuda_compile) and getattr(shared.compiled_model_state, 'first_pass_vae', False) and (shared.opts.cuda_compile_backend == "openvino_fx" or shared.opts.cuda_compile_backend == "openvino"):
shared.compiled_model_state.first_pass_vae = False
if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_model, "vqgan"):
model.vqgan.apply(sd_models_utils.convert_to_faketensors)
@@ -163,7 +163,7 @@ def full_vae_decode(latents, model):
del model.vae.orig_dtype
# delete vae after OpenVINO compile
- if 'VAE' in shared.opts.cuda_compile and shared.compiled_model_state.first_pass_vae and (shared.opts.cuda_compile_backend == "openvino_fx" or shared.opts.cuda_compile_backend == "openvino"):
+ if ('VAE' in shared.opts.cuda_compile) and getattr(shared.compiled_model_state, 'first_pass_vae', False) and (shared.opts.cuda_compile_backend == "openvino_fx" or shared.opts.cuda_compile_backend == "openvino"):
shared.compiled_model_state.first_pass_vae = False
if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_model, "vae"):
model.vae.apply(sd_models_utils.convert_to_faketensors)
@@ -187,7 +187,7 @@ def full_vae_encode(image, model):
log_debug('Moving to CPU: model=UNet')
unet_device = model.unet.device
sd_models.move_model(model.unet, devices.cpu)
- if not shared.opts.diffusers_offload_mode == "sequential" and hasattr(model, 'vae'):
+ if shared.opts.diffusers_offload_mode != "sequential" and hasattr(model, 'vae'):
sd_models.move_model(model.vae, devices.device)
vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default"
log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
@@ -247,6 +247,8 @@ def vae_postprocess(tensor, model, output_type='np'):
if tensor.ndim == 6 and tensor.shape[1] == 1:
tensor = tensor.squeeze(0)
images = model.video_processor.postprocess_video(tensor, output_type='pil')
+ if isinstance(images, list) and len(images) > 0 and isinstance(images[0], list):
+ images = [frame for batch in images for frame in batch]
elif hasattr(model, 'image_processor'):
if tensor.ndim == 5 and tensor.shape[1] == 3: # Qwen Image
tensor = tensor[:, :, 0]
diff --git a/modules/progress.py b/modules/progress.py
index 05f50077e..2fcd4dc22 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -69,7 +69,7 @@ class InternalProgressResponse(BaseModel):
def api_progress(req: ProgressRequest):
- active = req.id_task == current_task
+ active = (req.id_task == current_task) or (req.id_task == 'undefined' and len(shared.state.job) > 0)
queued = req.id_task in pending_tasks
completed = req.id_task in finished_tasks
paused = shared.state.paused
@@ -92,7 +92,7 @@ def api_progress(req: ProgressRequest):
id_live_preview = -1
textinfo = "Queued..." if queued else "Waiting..."
- debug_log(f'Preview: job={shared.state.job} active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} job={shared.state.preview_job} elapsed={elapsed:.3f}')
+ debug_log(f'Progress: job="{shared.state.job}" active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} job={shared.state.preview_job} elapsed={elapsed:.3f}')
if active and (req.id_live_preview != -1):
have_image = shared.state.set_current_image()
diff --git a/modules/res4lyf/res_unified_scheduler.py b/modules/res4lyf/res_unified_scheduler.py
index 061517f10..2b1d1bf0a 100644
--- a/modules/res4lyf/res_unified_scheduler.py
+++ b/modules/res4lyf/res_unified_scheduler.py
@@ -304,7 +304,7 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
return SchedulerOutput(prev_sample=x_next)
# GET COEFFICIENTS
- b, h_val = self._get_coefficients(sigma, sigma_next)
+ b, _h_val = self._get_coefficients(sigma, sigma_next)
if len(b) == 1:
res = b[0] * x0
diff --git a/modules/schedulers/perflow/scheduler_perflow.py b/modules/schedulers/perflow/scheduler_perflow.py
index e3a50feaf..3dcacc086 100644
--- a/modules/schedulers/perflow/scheduler_perflow.py
+++ b/modules/schedulers/perflow/scheduler_perflow.py
@@ -340,6 +340,27 @@ class PeRFlowScheduler(SchedulerMixin, ConfigMixin):
return PeRFlowSchedulerOutput(prev_sample=prev_sample, pred_original_sample=None)
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device)
+ else:
+ timestep = timestep.to(sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ if torch.is_floating_point(timestep):
+ timestep = timestep.round().to(dtype=torch.long)
+ return self.add_noise(sample, noise, timestep)
+
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
def add_noise(
self,
diff --git a/modules/schedulers/scheduler_bdia.py b/modules/schedulers/scheduler_bdia.py
index cffb5cb35..15b6b4caa 100644
--- a/modules/schedulers/scheduler_bdia.py
+++ b/modules/schedulers/scheduler_bdia.py
@@ -497,6 +497,26 @@ class BDIA_DDIMScheduler(SchedulerMixin, ConfigMixin):
return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device)
+ else:
+ timestep = timestep.to(sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ if torch.is_floating_point(timestep):
+ timestep = timestep.round().to(dtype=torch.long)
+ return self.add_noise(sample, noise, timestep)
+
def add_noise(
self,
original_samples: torch.Tensor,
diff --git a/modules/schedulers/scheduler_dc.py b/modules/schedulers/scheduler_dc.py
index 7121d4364..4a352777a 100644
--- a/modules/schedulers/scheduler_dc.py
+++ b/modules/schedulers/scheduler_dc.py
@@ -126,7 +126,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
use_karras_sigmas (`bool`, *optional*, defaults to `False`):
Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
- the sigmas are determined according to a sequence of noise levels {σi}.
+ the sigmas are determined according to a sequence of noise levels {sigma_i}.
timestep_spacing (`str`, defaults to `"linspace"`):
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
@@ -449,7 +449,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
model_output: torch.FloatTensor = None,
*args,
sample: torch.FloatTensor = None,
- order: int = None,
+ order: int | None = None,
**kwargs,
) -> torch.FloatTensor:
"""
@@ -488,13 +488,13 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
)
model_output_list = self.model_outputs
- s0 = self.timestep_list[-1]
+ self.timestep_list[-1]
m0 = model_output_list[-1]
assert m0 is not None
x = sample
if self.solver_p:
- raise NotImplementedError()
+ raise NotImplementedError
sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
@@ -534,7 +534,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
elif self.config.solver_type == "bh2":
B_h = torch.expm1(hh)
else:
- raise NotImplementedError()
+ raise NotImplementedError
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
@@ -579,7 +579,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
*args,
last_sample: torch.FloatTensor = None,
this_sample: torch.FloatTensor = None,
- order: int = None,
+ order: int | None = None,
**kwargs,
) -> torch.FloatTensor:
"""
@@ -669,7 +669,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
elif self.config.solver_type == "bh2":
B_h = torch.expm1(hh)
else:
- raise NotImplementedError()
+ raise NotImplementedError
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
@@ -811,7 +811,7 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
return loss
optimizer = torch.optim.AdamW([ratio_param], lr=0.1)
- for iter_ in range(self.num_iters):
+ for _ in range(self.num_iters):
optimizer.zero_grad()
loss = closure(ratio_param)
loss.backward()
@@ -1061,6 +1061,24 @@ class DCSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
"""
return sample
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device)
+ else:
+ timestep = timestep.to(sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ return self.add_noise(sample, noise, timestep)
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
def add_noise(
self,
diff --git a/modules/schedulers/scheduler_dpm_flowmatch.py b/modules/schedulers/scheduler_dpm_flowmatch.py
index 97472f754..3f6c6827c 100644
--- a/modules/schedulers/scheduler_dpm_flowmatch.py
+++ b/modules/schedulers/scheduler_dpm_flowmatch.py
@@ -234,7 +234,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
def set_timesteps(self,
- num_inference_steps: int = None,
+ num_inference_steps: int | None = None,
device: Union[str, torch.device] = None,
sigmas: Optional[List[float]] = None,
mu: Optional[float] = None,
@@ -295,6 +295,10 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
else:
num_inference_steps = len(sigmas)
self.num_inference_steps = num_inference_steps
+ if isinstance(sigmas, torch.Tensor):
+ sigmas = sigmas.detach().cpu().numpy()
+ else:
+ sigmas = np.asarray(sigmas, dtype=np.float64)
if self.config.sigma_schedule == "exponential":
if self.use_beta_sigmas:
diff --git a/modules/schedulers/scheduler_ersde.py b/modules/schedulers/scheduler_ersde.py
index a81fe14df..d42cfb38d 100644
--- a/modules/schedulers/scheduler_ersde.py
+++ b/modules/schedulers/scheduler_ersde.py
@@ -226,7 +226,11 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
def set_timesteps(self, num_inference_steps: Optional[int] = None, device: Union[str, torch.device] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, mu: Optional[float] = None):
if sigmas is not None:
# Flow-matching path: sigmas provided externally
- sigmas = np.array(sigmas, dtype=np.float64) if not isinstance(sigmas, np.ndarray) else sigmas.astype(np.float64)
+ if isinstance(sigmas, torch.Tensor):
+ sigmas = sigmas.detach().cpu().numpy()
+ elif not isinstance(sigmas, np.ndarray):
+ sigmas = np.asarray(sigmas, dtype=np.float64)
+ sigmas = sigmas.astype(np.float64, copy=False)
self.num_inference_steps = len(sigmas)
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device)
self._setup_flow(sigmas, device, mu)
diff --git a/modules/schedulers/scheduler_flashflow.py b/modules/schedulers/scheduler_flashflow.py
index e9a82c952..7d42b4edd 100644
--- a/modules/schedulers/scheduler_flashflow.py
+++ b/modules/schedulers/scheduler_flashflow.py
@@ -184,8 +184,8 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
def set_timesteps(
self,
- num_inference_steps: int = None,
- device: Union[str, torch.device] = None,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
sigmas: Optional[List[float]] = None,
mu: Optional[float] = None,
):
@@ -208,7 +208,11 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
sigmas = timesteps / self.config.num_train_timesteps
else:
- sigmas = np.array(sigmas).astype(np.float32)
+ if isinstance(sigmas, torch.Tensor):
+ sigmas = sigmas.detach().cpu().numpy()
+ else:
+ sigmas = np.asarray(sigmas, dtype=np.float32)
+ sigmas = sigmas.astype(np.float32, copy=False)
num_inference_steps = len(sigmas)
self.num_inference_steps = num_inference_steps
@@ -288,6 +292,7 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
+ noise_clip_std: float = 0.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]:
@@ -352,7 +357,9 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
device=model_output.device,
dtype=denoised.dtype,
)
- sample = sigma_next * noise + (1.0 - sigma_next) * denoised
+ if noise_clip_std > 0.0:
+ noise = noise.clamp(-noise_clip_std, noise_clip_std)
+ sample = sigma_next * s_noise * noise + (1.0 - sigma_next) * denoised
self._step_index += 1
sample = sample.to(model_output.dtype)
diff --git a/modules/schedulers/scheduler_tcd.py b/modules/schedulers/scheduler_tcd.py
index 83099217d..772602fc1 100644
--- a/modules/schedulers/scheduler_tcd.py
+++ b/modules/schedulers/scheduler_tcd.py
@@ -594,6 +594,26 @@ class TCDScheduler(SchedulerMixin, ConfigMixin):
return TCDSchedulerOutput(prev_sample=prev_sample, pred_noised_sample=pred_noised_sample)
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device)
+ else:
+ timestep = timestep.to(sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ if torch.is_floating_point(timestep):
+ timestep = timestep.round().to(dtype=torch.long)
+ return self.add_noise(sample, noise, timestep)
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
def add_noise(
self,
diff --git a/modules/schedulers/scheduler_tdd.py b/modules/schedulers/scheduler_tdd.py
index 05b49f35c..f7939f012 100644
--- a/modules/schedulers/scheduler_tdd.py
+++ b/modules/schedulers/scheduler_tdd.py
@@ -395,7 +395,7 @@ class TDDScheduler(DPMSolverSinglestepScheduler):
model_output_list: List[torch.FloatTensor],
*args,
sample: torch.FloatTensor = None,
- order: int = None,
+ order: int | None = None,
**kwargs,
) -> torch.FloatTensor:
timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
diff --git a/modules/schedulers/scheduler_ufogen.py b/modules/schedulers/scheduler_ufogen.py
index ff5f27eb3..908d9f95d 100644
--- a/modules/schedulers/scheduler_ufogen.py
+++ b/modules/schedulers/scheduler_ufogen.py
@@ -458,6 +458,26 @@ class UFOGenScheduler(SchedulerMixin, ConfigMixin):
return UFOGenSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device)
+ else:
+ timestep = timestep.to(sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ if torch.is_floating_point(timestep):
+ timestep = timestep.round().to(dtype=torch.long)
+ return self.add_noise(sample, noise, timestep)
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
def add_noise(
self,
diff --git a/modules/schedulers/scheduler_unipc_flowmatch.py b/modules/schedulers/scheduler_unipc_flowmatch.py
index bea747373..f94c4cab2 100644
--- a/modules/schedulers/scheduler_unipc_flowmatch.py
+++ b/modules/schedulers/scheduler_unipc_flowmatch.py
@@ -54,7 +54,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
use_karras_sigmas (`bool`, *optional*, defaults to `False`):
Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
- the sigmas are determined according to a sequence of noise levels {σi}.
+ the sigmas are determined according to a sequence of noise levels {sigma_i}.
use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
timestep_spacing (`str`, defaults to `"linspace"`):
@@ -181,6 +181,10 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
sigmas = np.linspace(self.sigma_max, self.sigma_min,
num_inference_steps +
1).copy()[:-1] # pyright: ignore
+ elif isinstance(sigmas, torch.Tensor):
+ sigmas = sigmas.detach().cpu().numpy()
+ else:
+ sigmas = np.asarray(sigmas, dtype=np.float32)
if self.config.use_dynamic_shifting:
sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore
@@ -311,7 +315,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
)
sigma = self.sigmas[self.step_index]
- alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+ _alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
if self.predict_x0:
if self.config.prediction_type == "flow_prediction":
@@ -350,7 +354,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
model_output: torch.Tensor,
*args,
sample: torch.Tensor = None,
- order: int = None, # pyright: ignore
+ order: int | None = None, # pyright: ignore
**kwargs,
) -> torch.Tensor:
"""
@@ -439,7 +443,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
elif self.config.solver_type == "bh2":
B_h = torch.expm1(hh)
else:
- raise NotImplementedError()
+ raise NotImplementedError
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
@@ -487,7 +491,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
*args,
last_sample: torch.Tensor = None,
this_sample: torch.Tensor = None,
- order: int = None, # pyright: ignore
+ order: int | None = None, # pyright: ignore
**kwargs,
) -> torch.Tensor:
"""
@@ -582,7 +586,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
elif self.config.solver_type == "bh2":
B_h = torch.expm1(hh)
else:
- raise NotImplementedError()
+ raise NotImplementedError
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
@@ -654,8 +658,14 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
model_output: torch.Tensor,
timestep: Union[int, torch.Tensor],
sample: torch.Tensor,
+ s_churn: float = 0.0,
+ s_tmin: float = 0.0,
+ s_tmax: float = float("inf"),
+ s_noise: float = 1.0,
+ noise_clip_std: float = 0.0,
return_dict: bool = True,
- generator=None) -> Union[SchedulerOutput, Tuple]:
+ generator=None,
+ **kwargs) -> Union[SchedulerOutput, Tuple]:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep UniPC.
@@ -752,6 +762,41 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
"""
return sample
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ """Forward process in flow-matching."""
+ sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+ if sample.device.type == "mps" and torch.is_floating_point(timestep):
+ # mps does not support float64
+ schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+ timestep = timestep.to(sample.device, dtype=torch.float32)
+ else:
+ schedule_timesteps = self.timesteps.to(sample.device)
+ timestep = timestep.to(sample.device)
+
+ if self.begin_index is None:
+ step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+ elif self.step_index is not None:
+ # add_noise is called after first denoising step (for inpainting)
+ step_indices = [self.step_index] * timestep.shape[0]
+ else:
+ # add noise is called before first denoising step to create initial latent(img2img)
+ step_indices = [self.begin_index] * timestep.shape[0]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < len(sample.shape):
+ sigma = sigma.unsqueeze(-1)
+
+ if noise is None:
+ noise = torch.randn_like(sample)
+
+ return sigma * noise + (1.0 - sigma) * sample
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
def add_noise(
self,
diff --git a/modules/schedulers/scheduler_vdm.py b/modules/schedulers/scheduler_vdm.py
index 35aab6e41..64885ee30 100644
--- a/modules/schedulers/scheduler_vdm.py
+++ b/modules/schedulers/scheduler_vdm.py
@@ -386,6 +386,24 @@ class VDMScheduler(SchedulerMixin, ConfigMixin):
return VDMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
+ def scale_noise(
+ self,
+ sample: torch.Tensor,
+ timestep: Union[float, torch.Tensor],
+ noise: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ if noise is None:
+ noise = torch.randn_like(sample)
+ if not isinstance(timestep, torch.Tensor):
+ timestep = torch.tensor([timestep], device=sample.device, dtype=sample.dtype)
+ else:
+ timestep = timestep.to(device=sample.device)
+ if timestep.ndim == 0:
+ timestep = timestep.unsqueeze(0)
+ if timestep.shape[0] != sample.shape[0]:
+ timestep = timestep.repeat(sample.shape[0])
+ return self.add_noise(sample, noise, timestep)
+
def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
"""
Adds noise to the original samples according to the noise schedule and the specified timesteps.
diff --git a/modules/scripts.py b/modules/scripts.py
index 941fa9bc1..daca52161 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -1,6 +1,6 @@
# compatibility with extensions that import scripts directly
from modules import scripts_manager
-from modules.scripts_manager import * # noqa: F403 # pylint: disable=wildcard-import
+from modules.scripts_manager import * # pylint: disable=wildcard-import
scripts_txt2img = None
diff --git a/modules/scripts_manager.py b/modules/scripts_manager.py
index 25bf6cf2e..57824b159 100644
--- a/modules/scripts_manager.py
+++ b/modules/scripts_manager.py
@@ -380,6 +380,11 @@ class ScriptRunner:
script.is_img2img = is_img2img
if path.startswith(paths.extensions_dir) and not path.startswith(paths.extensions_builtin_dir):
script.external = True
+ if is_control and script.external:
+ title = script.title()
+ if title not in control_extensions:
+ log.debug(f'Script: fn="{script.filename}" type=control title="{title}" skip')
+ return
if is_control: # this is messy but show is a legacy function that is not aware of control tab
v1 = script.show(script.is_txt2img)
v2 = script.show(script.is_img2img)
diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py
index 086010b9b..d840cba9a 100644
--- a/modules/sd_checkpoint.py
+++ b/modules/sd_checkpoint.py
@@ -24,7 +24,7 @@ warn_once = False
class CheckpointInfo:
- def __init__(self, filename, name=None, sha=None, subfolder=None, model_type: str = 'checkpoint'):
+ def __init__(self, filename, name=None, sha=None, subfolder=None, model_type: str = 'checkpoint', folder: str|None = None):
self.name = name
self.hash = sha
self.filename = filename
@@ -45,6 +45,8 @@ class CheckpointInfo:
relname = rel(filename, shared.opts.ckpt_dir)
elif relname.startswith(shared.opts.diffusers_dir):
relname = rel(filename, shared.opts.diffusers_dir)
+ elif relname.startswith(shared.opts.hfcache_dir):
+ relname = rel(filename, shared.opts.hfcache_dir)
elif relname.startswith(model_path):
relname = rel(filename, model_path)
elif relname.startswith(paths.script_path):
@@ -86,7 +88,7 @@ class CheckpointInfo:
self.shorthash = self.sha256[0:10] if self.sha256 else None
self.title = self.name if self.shorthash is None else f'{self.name} [{self.shorthash}]'
- self.path = self.filename
+ self.path = folder or self.filename
self.model_name = os.path.basename(self.name)
self.metadata = read_metadata_from_safetensors(filename)
# log.debug(f'Checkpoint: type={self.type} name={self.name} filename={self.filename} hash={self.shorthash} title={self.title}')
@@ -144,7 +146,7 @@ def list_models():
checkpoint_info.register()
diffusers_list = []
for repo in modelloader.load_diffusers_models(clear=True):
- checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash'])
+ checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash'], folder=repo['folder'])
diffusers_list.append(checkpoint_info)
if checkpoint_info.name is not None:
checkpoint_info.register()
diff --git a/modules/sd_detect.py b/modules/sd_detect.py
index 6cbbc0ca4..a3978f002 100644
--- a/modules/sd_detect.py
+++ b/modules/sd_detect.py
@@ -87,6 +87,10 @@ def guess_by_name(fn, current_guess):
new_guess = 'OmniGen'
elif 'sd3' in fn.lower():
new_guess = 'Stable Diffusion 3'
+ elif 'hidream-o1' in fn.lower():
+ new_guess = 'HiDreamO1'
+ elif 'hidream' in fn.lower():
+ new_guess = 'HiDream'
elif 'hidream' in fn.lower():
new_guess = 'HiDream'
elif 'zeta-chroma' in fn.lower() or 'zetachroma' in fn.lower():
@@ -97,6 +101,8 @@ def guess_by_name(fn, current_guess):
new_guess = 'FLUX2 Klein'
elif 'flux.2' in fn.lower():
new_guess = 'FLUX2'
+ elif 'ultraflux' in fn.lower():
+ new_guess = 'UltraFlux'
elif 'flux' in fn.lower() or 'flex.1' in fn.lower():
size = round(os.path.getsize(fn) / 1024 / 1024) if os.path.isfile(fn) else 0
if size > 11000 and size < 16000:
@@ -126,7 +132,7 @@ def guess_by_name(fn, current_guess):
new_guess = 'Kandinsky 2.2'
elif 'kandinsky-3' in fn.lower():
new_guess = 'Kandinsky 3.0'
- elif 'kandinsky-5.0' in fn.lower() and '2i' not in fn.lower():
+ elif 'kandinsky-5.0' in fn.lower():
new_guess = 'Kandinsky 5.0'
elif 'hunyuanimage3' in fn.lower() or 'hunyuanimage-3' in fn.lower():
new_guess = 'HunyuanImage3'
@@ -160,20 +166,24 @@ def guess_by_name(fn, current_guess):
new_guess = 'Step1X-Edit'
elif 'vibe-image-edit' in fn.lower():
new_guess = 'VIBE'
+ elif 'joyai-image-edit' in fn.lower() or 'joy-image-edit' in fn.lower():
+ new_guess = 'Joy'
if debug_load:
log.trace(f'Autodetect: method=name file="{fn}" previous="{current_guess}" current="{new_guess}"')
return new_guess or current_guess
def guess_by_diffusers(fn, current_guess):
- exclude_by_name = ['ostris/Flex.2-preview'] # pipeline may be misleading
+ exclude_by_name = ['ostris/Flex.2-preview', 'Owen777/UltraFlux-v1', './pretrain/FLUX.1-dev'] # pipeline may be misleading
if not os.path.isdir(fn):
return current_guess, None
index = os.path.join(fn, 'model_index.json')
if os.path.exists(index) and os.path.isfile(index):
index = shared.readfile(index, silent=True, as_type="dict")
name = index.get('_name_or_path', None)
- if name is not None and name in exclude_by_name:
+ if debug_load:
+ log.trace(f'Autodetect: method=diffusers file="{fn}" name="{name}"')
+ if (name is not None) and (name in exclude_by_name):
return current_guess, None
cls = index.get('_class_name', None)
if isinstance(cls, list):
@@ -242,9 +252,17 @@ def detect_pipeline(f: str, op: str = 'model'):
try:
guess = 'Stable Diffusion XL' if ('XL' in f.upper() or 'SDNQ' in f.upper()) else 'Stable Diffusion' # set default guess
guess = guess_by_size(f, guess)
+ if debug_load:
+ log.trace(f'Autodetect: type=size guess="{guess}" file="{f}"')
guess = guess_by_name(f, guess)
+ if debug_load:
+ log.trace(f'Autodetect: type=name guess="{guess}" file="{f}"')
guess, pipeline = guess_by_diffusers(f, guess)
+ if debug_load:
+ log.trace(f'Autodetect: type=diffusers guess="{guess}" file="{f}"')
guess = guess_variant(f, guess)
+ if debug_load:
+ log.trace(f'Autodetect: type=variant guess="{guess}" file="{f}"')
pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
log.info(f'Autodetect {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}"')
if debug_load is not None:
diff --git a/modules/sd_hijack_hfhub.py b/modules/sd_hijack_hfhub.py
index 46d0a58d9..17be5a172 100644
--- a/modules/sd_hijack_hfhub.py
+++ b/modules/sd_hijack_hfhub.py
@@ -16,7 +16,7 @@ def http_get_hijack(*args, **kwargs):
fn = kwargs.get("displayed_filename", None)
size = kwargs.get("expected_size", None)
if fn and not fn.endswith(".json") and size is not None and size > 10240:
- log.debug(f'Download start: type=http fn="{fn}" size={size}')
+ log.debug(f'Download: type=http fn="{fn}" size={size}')
debug(f'Download start: type=http args={args} kwargs={kwargs}')
t0 = time.time()
res = orig_http_get(*args, **kwargs)
@@ -34,7 +34,7 @@ def xet_get_hijack(*args, **kwargs):
fn = kwargs.get("displayed_filename", None)
size = kwargs.get("expected_size", None)
if fn and not fn.endswith(".json"):
- log.debug(f'Download start: type=xet fn="{fn}" size={size}')
+ log.debug(f'Download: type=xet fn="{fn}" size={size}')
debug(f'Download start: type=xet args={args} kwargs={kwargs}')
res = orig_xet_get(*args, **kwargs)
debug(f'Download end: type=xet res={res}')
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 296f32d5b..288d10ab3 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -51,13 +51,16 @@ pipe_switch_task_exclude = [
'NucleusMoEImagePipeline',
'AuraFlowPipeline',
'ChronoEditPipeline',
+ 'Kandinsky5I2IPipeline',
'GoogleNanoBananaPipeline',
+ 'Step1XEditPipeline',
]
i2i_pipes = [
'LEditsPPPipelineStableDiffusion', 'LEditsPPPipelineStableDiffusionXL',
'OmniGenPipeline', 'OmniGen2Pipeline',
'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
+ 'Step1XEditPipeline',
]
@@ -374,6 +377,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
from pipelines.model_auraflow import load_auraflow
sd_model = load_auraflow(checkpoint_info, diffusers_load_config)
allow_post_quant = False
+ elif model_type in ['UltraFlux']:
+ from pipelines.model_ultraflux import load_ultraflux
+ sd_model = load_ultraflux(checkpoint_info, diffusers_load_config)
+ allow_post_quant = False
elif model_type in ['FLUX']:
from pipelines.model_flux import load_flux
sd_model = load_flux(checkpoint_info, diffusers_load_config)
@@ -426,6 +433,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
from pipelines.model_omnigen import load_omnigen
sd_model = load_omnigen(checkpoint_info, diffusers_load_config)
allow_post_quant = False
+ elif model_type in ['HiDreamO1']:
+ from pipelines.model_hidream import load_hidream_o1
+ sd_model = load_hidream_o1(checkpoint_info, diffusers_load_config)
+ allow_post_quant = False
elif model_type in ['HiDream']:
from pipelines.model_hidream import load_hidream
sd_model = load_hidream(checkpoint_info, diffusers_load_config)
@@ -462,6 +473,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
from pipelines.model_vibe import load_vibe
sd_model = load_vibe(checkpoint_info, diffusers_load_config)
allow_post_quant = False
+ elif model_type in ['Joy']:
+ from pipelines.model_joy import load_joy
+ sd_model = load_joy(checkpoint_info, diffusers_load_config)
+ allow_post_quant = False
elif model_type in ['Qwen']:
from pipelines.model_qwen import load_qwen
sd_model = load_qwen(checkpoint_info, diffusers_load_config)
@@ -482,7 +497,7 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
from pipelines.model_kandinsky import load_kandinsky3
sd_model = load_kandinsky3(checkpoint_info, diffusers_load_config)
allow_post_quant = False
- elif model_type in ['Kandinsky 5.0'] and model_type:
+ elif model_type in ['Kandinsky 5.0']:
from pipelines.model_kandinsky import load_kandinsky5
sd_model = load_kandinsky5(checkpoint_info, diffusers_load_config)
allow_post_quant = False
@@ -963,17 +978,39 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di
move_model(sd_model, devices.device)
timer.load.record("move")
- if shared.opts.ipex_optimize:
- sd_model = sd_models_compile.ipex_optimize(sd_model)
-
- if ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none'):
- sd_model = sd_models_compile.compile_diffusers(sd_model)
- timer.load.record("compile")
-
except Exception as e:
log.error(f"Load {op}: {e}")
errors.display(e, "Model")
+ try:
+ if shared.opts.ipex_optimize:
+ sd_model = sd_models_compile.ipex_optimize(sd_model)
+
+ if (shared.opts.cuda_compile_backend != 'none') and len(shared.opts.cuda_compile) > 0:
+ if 'components' in shared.opts.cuda_compile_options:
+ sd_model = sd_models_compile.compile_diffusers(sd_model, apply_to_components=True)
+ else:
+ if 'Model' in shared.opts.cuda_compile:
+ if hasattr(sd_model, "unet"):
+ sd_model.unet = sd_models_compile.compile_diffusers(sd_model.unet, apply_to_components=False)
+ if hasattr(sd_model, "transformer"):
+ sd_model.transformer = sd_models_compile.compile_diffusers(sd_model.transformer, apply_to_components=False)
+ if 'TE' in shared.opts.cuda_compile:
+ if hasattr(sd_model, "text_encoder"):
+ sd_model.text_encoder = sd_models_compile.compile_diffusers(sd_model.text_encoder, apply_to_components=False)
+ if hasattr(sd_model, "text_encoder_2"):
+ sd_model.text_encoder_2 = sd_models_compile.compile_diffusers(sd_model.text_encoder_2, apply_to_components=False)
+ if hasattr(sd_model, "text_encoder_3"):
+ sd_model.text_encoder_3 = sd_models_compile.compile_diffusers(sd_model.text_encoder_3, apply_to_components=False)
+ if 'VAE' in shared.opts.cuda_compile:
+ if hasattr(sd_model, "vae"):
+ sd_model.vae = sd_models_compile.compile_diffusers(sd_model.vae, apply_to_components=False)
+
+ timer.load.record("compile")
+ except Exception as e:
+ log.error(f"Compile {op}: {e}")
+ errors.display(e, "Compile")
+
if shared.opts.diffusers_offload_mode != 'balanced':
devices.torch_gc(force=True, reason='load')
if sd_model is not None:
@@ -1252,6 +1289,8 @@ def set_diffuser_pipe(pipe, new_pipe_type):
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
log.trace(f"Pipeline class change requested: target={new_pipe_type} fn={fn}") # pylint: disable=protected-access
log.warning(f'Pipeline class change failed: type={new_pipe_type} pipeline={cls} {e}')
+ if debug_load:
+ errors.display(e, 'Pipeline switch')
has_errors = True
if not hasattr(pipe, 'config') or has_errors:
try: # maybe a wrapper pipeline so just change the class
@@ -1269,6 +1308,8 @@ def set_diffuser_pipe(pipe, new_pipe_type):
return pipe
except Exception as e: # pylint: disable=unused-variable
log.warning(f'Pipeline class set failed: type={new_pipe_type} pipeline={cls} {e}')
+ if debug_load:
+ errors.display(e, 'Pipeline switch')
has_errors = True
return pipe
@@ -1290,7 +1331,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
add_noise_pred_to_diffusers_callback(new_pipe.pipe)
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
- log.debug(f"Pipeline class change: original={cls} target={new_pipe.__class__.__name__} device={pipe.device} fn={fn}") # pylint: disable=protected-access
+ log.debug(f"Pipeline class change: source={cls} target={new_pipe.__class__.__name__} device={pipe.device} fn={fn}") # pylint: disable=protected-access
if shared.opts.diffusers_offload_mode == 'none':
move_model(new_pipe, pipe.device)
@@ -1482,3 +1523,13 @@ def save_model(name: str, path: str | None = None, shard: str = "5GB", overwrite
log.error(f'Save model: path="{model_name}" {e}')
errors.display(e, 'Save model')
return f'Error: {e}'
+
+
+def list_hfcache():
+ checkpoints = []
+ for f in os.scandir(shared.opts.hfcache_dir):
+ if not os.path.isdir(f) or not f.name.startswith('models--'):
+ continue
+ checkpoint = CheckpointInfo(filename=f.path, name=path_to_repo(f.name), model_type='hfcache')
+ checkpoints.append(checkpoint)
+ return checkpoints
diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py
index 7c0851434..411b2f912 100644
--- a/modules/sd_models_compile.py
+++ b/modules/sd_models_compile.py
@@ -94,7 +94,7 @@ def compile_onediff(sd_model):
log.warning(f"Model compile: task=onediff {e}")
return sd_model
- debug_log(f"Model compile: task=onediff pipeline={sd_model.__class__.__name__}")
+ log.info(f"Model compile: task=onediff pipeline={sd_model.__class__.__name__} precompile={'precompile' in shared.opts.cuda_compile_options}")
try:
t0 = time.time()
# For some reason compiling the text_encoder, when it is used by
@@ -144,7 +144,7 @@ def compile_stablefast(sd_model):
# config.trace_scheduler = False
# config.enable_cnn_optimization
# config.prefer_lowp_gemm
- debug_log(f"Model compile: task=stablefast config={config.__dict__}")
+ log.debug(f"Model compile: task=stablefast config={config.__dict__} precompile={'precompile' in shared.opts.cuda_compile_options} {config}")
try:
t0 = time.time()
sd_model = sf.compile(sd_model, config)
@@ -165,14 +165,18 @@ def compile_torch(sd_model, apply_to_components=True, op="Model"):
t0 = time.time()
import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
torch._dynamo.reset() # pylint: disable=protected-access
- log.debug(f"{op} compile: task=torch backends={torch._dynamo.list_backends()}") # pylint: disable=protected-access
- debug_log(f"{op} compile: options={shared.opts.cuda_compile_options} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} targets={shared.opts.cuda_compile}")
+ log.debug(f"{op} compile: task=torch available={torch._dynamo.list_backends()}") # pylint: disable=protected-access
+ is_repeated = hasattr(sd_model, 'compile_repeated_blocks') and 'repeated' in shared.opts.cuda_compile_options and not sd_model.__class__.__name__.startswith("Autoencoder")
+ log.debug(f"{op} compile: options={shared.opts.cuda_compile_options} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} repeated={is_repeated} components={apply_to_components} targets={shared.opts.cuda_compile}")
compiled_components = []
def torch_compile_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
- compiled_components.append(model.__class__.__name__)
- if hasattr(model, 'compile_repeated_blocks') and 'repeated' in shared.opts.cuda_compile_options:
+ setup_logging() # compile messes with logging so reset is needed
+ log.debug(f"Compile: cls={sd_model.__class__.__name__} apply")
+ name = model.__class__.__name__ if callable(model) else model.__name__
+ compiled_components.append(name)
+ if is_repeated:
model.compile_repeated_blocks(
mode=shared.opts.cuda_compile_mode,
backend=shared.opts.cuda_compile_backend,
@@ -194,7 +198,8 @@ def compile_torch(sd_model, apply_to_components=True, op="Model"):
fullgraph='fullgraph' in shared.opts.cuda_compile_options,
dynamic='dynamic' in shared.opts.cuda_compile_options,
)
- devices.torch_gc()
+ devices.torch_gc(force=True, reason='compile')
+ setup_logging() # compile messes with logging so reset is needed
return model
if shared.opts.cuda_compile_backend == "openvino_fx" or shared.opts.cuda_compile_backend == "openvino":
@@ -207,11 +212,17 @@ def compile_torch(sd_model, apply_to_components=True, op="Model"):
pass # pylint: disable=unused-import
verbose = debug or 'verbose' in shared.opts.cuda_compile_options
log_level = logging.WARNING if verbose else logging.CRITICAL # pylint: disable=protected-access
+
+ # configure torch.dynamo
if hasattr(torch, '_logging'):
torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
+ setup_logging() # dynamo messes with logging so reset is needed
torch._dynamo.config.verbose = verbose # pylint: disable=protected-access
torch._dynamo.config.suppress_errors = not verbose # pylint: disable=protected-access
+ if 'dynamic' in shared.opts.cuda_compile_options:
+ torch._dynamo.config.capture_dynamic_output_shape_ops = True # pylint: disable=protected-access
+ # configure torch.inductor
try:
torch._inductor.config.conv_1x1_as_mm = True # pylint: disable=protected-access
torch._inductor.config.coordinate_descent_tuning = True # pylint: disable=protected-access
@@ -227,16 +238,14 @@ def compile_torch(sd_model, apply_to_components=True, op="Model"):
else:
sd_model = torch_compile_model(sd_model, op=op)
- setup_logging() # compile messes with logging so reset is needed
if apply_to_components and 'precompile' in shared.opts.cuda_compile_options:
try:
- log.debug(f"{op} compile: task=torch precompile")
+ log.debug(f"{op} compile: precompile start")
sd_model("dummy prompt")
except Exception:
pass
t1 = time.time()
- log.info(f"{op} compile: task=torch time={t1-t0:.2f}")
- debug_log(f"{op} compile: task=torch completed components={compiled_components} targets={shared.opts.cuda_compile} verbose={verbose} time={t1-t0:.2f}")
+ log.info(f"{op} compile: task=torch components={compiled_components} time={t1-t0:.2f}")
except Exception as e:
log.warning(f"{op} compile: task=torch {e}")
errors.display(e, 'Compile')
@@ -261,7 +270,7 @@ def compile_deepcache(sd_model):
except Exception as e:
log.warning(f'Model compile: task=deepcache {e}')
return sd_model
- debug_log(f"Model compile: task=deepcache pipeline={sd_model.__class__.__name__} interval={shared.opts.deep_cache_interval}")
+ log.debug(f"Model compile: task=deepcache pipeline={sd_model.__class__.__name__} interval={shared.opts.deep_cache_interval}")
t0 = time.time()
check_deepcache(False)
deepcache_worker = DeepCacheSDHelper(pipe=sd_model)
@@ -276,7 +285,7 @@ def compile_diffusers(sd_model, apply_to_components=True, op="Model"):
if shared.opts.cuda_compile_backend == 'none':
log.warning(f'{op} compile enabled but no backend specified')
return sd_model
- log.info(f"{op} compile: pipeline={sd_model.__class__.__name__} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} options={shared.opts.cuda_compile_options} compile={shared.opts.cuda_compile}")
+ log.info(f"{op} compile: pipeline={sd_model.__class__.__name__} backend={shared.opts.cuda_compile_backend} options={shared.opts.cuda_compile_options}")
if shared.opts.cuda_compile_backend == 'onediff':
sd_model = compile_onediff(sd_model)
elif shared.opts.cuda_compile_backend == 'stable-fast':
diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py
index e6f5d10e8..2a9db5f79 100644
--- a/modules/sd_models_utils.py
+++ b/modules/sd_models_utils.py
@@ -73,6 +73,18 @@ def path_to_repo(checkpoint_info):
return repo_id
+def repo_to_path(repo_id):
+ if repo_id.name.startswith('Diffusers'):
+ folder = repo_id[len('Diffusers'):].lstrip('/')
+ folder = 'models--' + repo_id.replace('/', '--')
+ folder = os.path.join(shared.opts.diffusers_dir, repo_id)
+ if os.path.exists(folder):
+ return folder
+ if os.path.exists(repo_id.filename):
+ return repo_id.filename
+ return ''
+
+
def convert_to_faketensors(tensor):
try:
fake_module = torch._subclasses.fake_tensor.FakeTensorMode(allow_non_fake_inputs=True) # pylint: disable=protected-access
@@ -222,10 +234,18 @@ def apply_function_to_model(sd_model, function, options, op=None):
if "VAE" in options:
if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'):
if op == "compile":
- sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model)
- sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model)
+ if hasattr(sd_model.vae, 'decoder'):
+ sd_model.vae.decoder = function(sd_model.vae.decoder, op="vae_decoder", sd_model=sd_model)
+ else:
+ sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model)
else:
sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model)
+ if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'encode'):
+ if op == "compile":
+ if hasattr(sd_model.vae, 'encoder'):
+ sd_model.vae.encoder = function(sd_model.vae.encoder, op="vae_encoder", sd_model=sd_model)
+ else:
+ sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model)
if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'):
if op == "compile":
sd_model.movq.decode = function(sd_model.movq.decode, op="movq_decode", sd_model=sd_model)
diff --git a/modules/sd_offload.py b/modules/sd_offload.py
index da478b883..80ca7fe58 100644
--- a/modules/sd_offload.py
+++ b/modules/sd_offload.py
@@ -234,7 +234,6 @@ class OffloadHook(accelerate.hooks.ModelHook):
return False
return True
- @torch.compiler.disable
def pre_forward(self, module, *args, **kwargs):
_id = id(module)
@@ -277,13 +276,20 @@ class OffloadHook(accelerate.hooks.ModelHook):
log.trace(f'Offload: type=balanced op=dispatch map={device_map}')
if device_map is not None:
skip_keys = getattr(module, "_skip_keys", None)
- module = accelerate.dispatch_model(module,
- main_device=torch.device(devices.device),
- device_map=device_map,
- offload_dir=offload_dir,
- skip_keys=skip_keys,
- force_hooks=True,
- )
+ try:
+ module = accelerate.dispatch_model(module,
+ main_device=torch.device(devices.device),
+ device_map=device_map,
+ offload_dir=offload_dir,
+ skip_keys=skip_keys,
+ force_hooks=True,
+ )
+ except Exception as e: # reapply hook
+ log.warning(f'Offload: type=balanced op=dispatch module={module.__class__.__name__} {e}')
+ module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
+ module.balanced_offload_device_map = None
+ sd_models.move_model(module, devices.device, force=True)
+ module = accelerate.hooks.add_hook_to_module(module, self, append=True)
module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
module.balanced_offload_device_map = device_map
module.balanced_offload_max_memory = max_memory
@@ -298,7 +304,6 @@ class OffloadHook(accelerate.hooks.ModelHook):
self.last_pre = _id
return args, kwargs
- @torch.compiler.disable
def post_forward(self, module, output):
if self.last_post != id(module):
self.last_post = id(module)
@@ -500,7 +505,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc
continue
module.module_name = module_name
module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name)
- apply_balanced_offload_to_module(module, op='apply')
+ apply_balanced_offload_to_module(module, op='apply', force=force)
if not silent:
report_model_stats(module_name, module)
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 8bad35232..92e7ca12a 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -169,25 +169,38 @@ def load_vae(model_file, vae_file=None, vae_source="unknown-source"):
vae_config = sd_detect.get_load_config(model_file, model_type, config_type='json')
if vae_config is not None:
diffusers_load_config['config'] = os.path.join(vae_config, 'vae')
- log.info(f'Load module: type=VAE model="{vae_file}" source={vae_source} config={diffusers_load_config}')
try:
import diffusers
- if os.path.isfile(vae_file):
+ vae_class = None
+ vae_loader = None
+ if shared.sd_loaded and getattr(shared.sd_model, 'vae', None) is not None:
+ vae_class = shared.sd_model.vae.__class__
+ vae_loader = vae_class.from_single_file if os.path.isfile(vae_file) else vae_class.from_pretrained
+ elif os.path.isfile(vae_file):
if os.path.getsize(vae_file) > 1310944880: # 1.3GB
- vae = diffusers.ConsistencyDecoderVAE.from_pretrained('openai/consistency-decoder', **diffusers_load_config) # consistency decoder does not have from single file, so we'll just download it once more
+ vae_class = diffusers.ConsistencyDecoderVAE
+ vae_loader = vae_class.from_pretrained
+ vae_file = 'openai/consistency-decoder'
elif os.path.getsize(vae_file) < 10000000: # 10MB
- vae = diffusers.AutoencoderTiny.from_single_file(vae_file, **diffusers_load_config)
- else:
- vae = diffusers.AutoencoderKL.from_single_file(vae_file, **diffusers_load_config)
- if getattr(vae.config, 'scaling_factor', 0) == 0.18125 and shared.sd_model_type == 'sdxl':
- vae.config.scaling_factor = 0.13025
- log.debug('Setting model: component=VAE fix scaling factor')
- vae = vae.to(devices.dtype_vae)
+ vae_class = diffusers.AutoencoderTiny
+ vae_loader = vae_class.from_single_file
+ else: # fallback
+ vae_class = diffusers.AutoencoderKL
+ # if getattr(vae.config, 'scaling_factor', 0) == 0.18125 and shared.sd_model_type == 'sdxl':
+ # vae.config.scaling_factor = 0.13025
+ # log.debug('Setting model: component=VAE fix scaling factor')
+ vae_loader = vae_class.from_single_file
else:
if 'consistency-decoder' in vae_file:
- vae = diffusers.ConsistencyDecoderVAE.from_pretrained(vae_file, **diffusers_load_config)
- else:
- vae = diffusers.AutoencoderKL.from_pretrained(vae_file, **diffusers_load_config)
+ vae_class = diffusers.ConsistencyDecoderVAE
+ else: # fallback
+ vae_class = diffusers.AutoencoderKL
+ vae_loader = vae_class.from_pretrained
+ if vae_loader is not None:
+ log.info(f'Load module: type=VAE model="{vae_file}" source={vae_source} cls={vae_class.__name__} config={diffusers_load_config}')
+ vae = vae_loader(vae_file, **diffusers_load_config)
+ vae = vae.to(devices.dtype_vae)
+
global loaded_vae_file # pylint: disable=global-statement
loaded_vae_file = os.path.basename(vae_file)
# log.debug(f'Diffusers VAE config: {vae.config}')
diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py
index 4fe6b10c4..f383fe5b7 100644
--- a/modules/sdnq/common.py
+++ b/modules/sdnq/common.py
@@ -513,6 +513,10 @@ module_skip_keys_dict = {
[".emb_in", ".txt_in", ".vid_in", ".emb_scale", ".vid_out", ".vid_out_norm", ".vid_out_ada"],
{}
],
+ "HiDreamO1Qwen3VLTransformer": [
+ ["lm_head", "embed_tokens", "x_embedder", "t_embedder1", "final_layer2", "patch_embed", "pos_embed"],
+ {}
+ ],
}
module_skip_keys_dict["LongCatImageTransformer2DModel"] = module_skip_keys_dict["FluxTransformer2DModel"]
diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py
index 2a0b08546..886d4e571 100644
--- a/modules/sdnq/layers/linear/linear_int8.py
+++ b/modules/sdnq/layers/linear/linear_int8.py
@@ -50,7 +50,17 @@ def int8_matmul(
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
- return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
+ dequantized_weight = self.sdnq_dequantizer(
+ self.weight,
+ self.scale,
+ self.zero_point,
+ self.svd_up,
+ self.svd_down,
+ skip_quantized_matmul=True,
+ )
+ if input.dtype != dequantized_weight.dtype:
+ input = input.to(dtype=dequantized_weight.dtype)
+ return torch.nn.functional.linear(input, dequantized_weight, self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py
index 6fb9f2d9b..fa7ce30dd 100644
--- a/modules/sdnq/loader.py
+++ b/modules/sdnq/loader.py
@@ -25,6 +25,31 @@ def unset_config_on_save(quantization_config: SDNQConfig) -> SDNQConfig:
return quantization_config
+def normalize_tied_weights_keys_for_save(model: ModelMixin, is_pipeline: bool = False) -> list[tuple[torch.nn.Module, object]]:
+ normalized_modules = []
+ modules_to_walk = []
+ if is_pipeline:
+ for module_name in get_module_names(model):
+ module = getattr(model, module_name, None)
+ if isinstance(module, torch.nn.Module):
+ modules_to_walk.append(module)
+ elif isinstance(model, torch.nn.Module):
+ modules_to_walk.append(model)
+
+ for root_module in modules_to_walk:
+ for submodule in root_module.modules():
+ tied_weights_keys = getattr(submodule, "_tied_weights_keys", None)
+ if isinstance(tied_weights_keys, list):
+ normalized_modules.append((submodule, tied_weights_keys))
+ submodule._tied_weights_keys = {key: key for key in tied_weights_keys} # pylint: disable=protected-access
+ return normalized_modules
+
+
+def restore_tied_weights_keys_after_save(normalized_modules: list[tuple[torch.nn.Module, object]]) -> None:
+ for submodule, tied_weights_keys in normalized_modules:
+ submodule._tied_weights_keys = tied_weights_keys # pylint: disable=protected-access
+
+
def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "5GB", is_pipeline: bool = False, sdnq_config: SDNQConfig | None = None) -> None:
if is_pipeline:
for module_name in get_module_names(model):
@@ -39,7 +64,11 @@ def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "5
if hasattr(model, "quantization_config") and isinstance(model.quantization_config, SDNQConfig):
model.quantization_config = unset_config_on_save(model.quantization_config)
- model.save_pretrained(model_path, max_shard_size=max_shard_size) # actual save
+ normalized_modules = normalize_tied_weights_keys_for_save(model, is_pipeline=is_pipeline)
+ try:
+ model.save_pretrained(model_path, max_shard_size=max_shard_size) # actual save
+ finally:
+ restore_tied_weights_keys_after_save(normalized_modules)
quantization_config_path = os.path.join(model_path, "quantization_config.json")
if sdnq_config is not None: # if provided, save global config
diff --git a/modules/sdnq/triton_mm.py b/modules/sdnq/triton_mm.py
index 955761e88..495d97945 100644
--- a/modules/sdnq/triton_mm.py
+++ b/modules/sdnq/triton_mm.py
@@ -12,31 +12,16 @@ import torch
import triton
import triton.language as tl
-try:
- from .common import is_rdna2_and_older
-except Exception:
- is_rdna2_and_older = False
-if is_rdna2_and_older:
- matmul_configs = [
- triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
- for BM in [64, 128]
- for BN in [64, 128]
- for BK in [64]
- for GM in [2, 4]
- for w in [2, 4]
- for s in [2]
- ]
-else:
- matmul_configs = [
- triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
- for BM in [32, 64, 128, 256]
- for BN in [32, 64, 128, 256]
- for BK in [32, 64, 128]
- for GM in [4, 8]
- for w in [4, 8]
- for s in [2]
- ]
+matmul_configs = [
+ triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
+ for BM in [64, 128, 256]
+ for BN in [64, 128, 256]
+ for BK in [64, 128]
+ for GM in [2, 4, 8]
+ for w in [2, 4, 8]
+ for s in [2]
+]
@triton.autotune(configs=matmul_configs, key=["M", "N", "K", "stride_bk", "ACCUMULATOR_DTYPE"], cache_results=True)
diff --git a/modules/seedvr/src/common/config.py b/modules/seedvr/src/common/config.py
index 5b8780bc7..c04997536 100644
--- a/modules/seedvr/src/common/config.py
+++ b/modules/seedvr/src/common/config.py
@@ -15,7 +15,7 @@ except Exception as e:
raise
-def load_config(path: str, argv: List[str] = None):
+def load_config(path: str, argv: List[str] | None = None):
"""
Load a configuration. Will resolve inheritance.
"""
diff --git a/modules/seedvr/src/common/diffusion/samplers/euler.py b/modules/seedvr/src/common/diffusion/samplers/euler.py
index 185669acb..d4892a3f1 100644
--- a/modules/seedvr/src/common/diffusion/samplers/euler.py
+++ b/modules/seedvr/src/common/diffusion/samplers/euler.py
@@ -27,6 +27,7 @@ from torch.nn import functional as F
from ..types import PredictionType
from ..utils import expand_dims
from .base import Sampler, SamplerModelArgs
+import itertools
class EulerSampler(Sampler):
@@ -46,9 +47,8 @@ class EulerSampler(Sampler):
# Optimisations VRAM
original_dtype = x.dtype
- device = x.device
- for t, s in zip(timesteps[:-1], timesteps[1:]):
+ for t, s in itertools.pairwise(timesteps):
# Appel du modèle avec monitoring
pred = f(SamplerModelArgs(x, t, i))
diff --git a/modules/seedvr/src/common/distributed/meta_init_utils.py b/modules/seedvr/src/common/distributed/meta_init_utils.py
index 9a6570646..d0e64a952 100644
--- a/modules/seedvr/src/common/distributed/meta_init_utils.py
+++ b/modules/seedvr/src/common/distributed/meta_init_utils.py
@@ -14,7 +14,7 @@
import torch
from torch import nn
-from ....rotary_embedding import RotaryEmbedding
+from modules.seedvr.rotary_embedding import RotaryEmbedding
__all__ = ["meta_non_persistent_buffer_init_fn"]
diff --git a/modules/seedvr/src/common/distributed/ops.py b/modules/seedvr/src/common/distributed/ops.py
index f3101b62c..1cbade014 100644
--- a/modules/seedvr/src/common/distributed/ops.py
+++ b/modules/seedvr/src/common/distributed/ops.py
@@ -22,7 +22,7 @@ import torch
import torch.distributed as dist
from torch import Tensor
-from ..cache import Cache
+from modules.seedvr.src.common.cache import Cache
from .advanced import (
get_sequence_parallel_group,
get_sequence_parallel_rank,
diff --git a/modules/seedvr/src/core/generation.py b/modules/seedvr/src/core/generation.py
index 401d8b874..e5b231404 100644
--- a/modules/seedvr/src/core/generation.py
+++ b/modules/seedvr/src/core/generation.py
@@ -1,10 +1,10 @@
import torch
from torchvision.transforms import Compose, Lambda, Normalize
-from ..optimization.performance import optimized_video_rearrange, optimized_single_video_rearrange, optimized_sample_to_image_format
-from ..common.seed import set_seed
-from ..data.image.transforms.divisible_crop import DivisibleCrop
-from ..data.image.transforms.na_resize import NaResize
-from ..utils.color_fix import wavelet_reconstruction
+from modules.seedvr.src.optimization.performance import optimized_video_rearrange, optimized_single_video_rearrange, optimized_sample_to_image_format
+from modules.seedvr.src.common.seed import set_seed
+from modules.seedvr.src.data.image.transforms.divisible_crop import DivisibleCrop
+from modules.seedvr.src.data.image.transforms.na_resize import NaResize
+from modules.seedvr.src.utils.color_fix import wavelet_reconstruction
diff --git a/modules/seedvr/src/core/infer.py b/modules/seedvr/src/core/infer.py
index 5be430316..0e4a24545 100644
--- a/modules/seedvr/src/core/infer.py
+++ b/modules/seedvr/src/core/infer.py
@@ -1,8 +1,8 @@
from typing import List, Optional, Tuple, Union
import torch
from einops import rearrange
-from ..common.diffusion import classifier_free_guidance_dispatcher, create_sampler_from_config, create_sampling_timesteps_from_config, create_schedule_from_config
-from ..models.dit_v2 import na
+from modules.seedvr.src.common.diffusion import classifier_free_guidance_dispatcher, create_sampler_from_config, create_sampling_timesteps_from_config, create_schedule_from_config
+from modules.seedvr.src.models.dit_v2 import na
def optimized_channels_to_last(tensor: torch.Tensor) -> torch.Tensor:
diff --git a/modules/seedvr/src/core/model_manager.py b/modules/seedvr/src/core/model_manager.py
index 3625af263..a3ad6db7b 100644
--- a/modules/seedvr/src/core/model_manager.py
+++ b/modules/seedvr/src/core/model_manager.py
@@ -2,9 +2,9 @@ import os
import torch
from safetensors.torch import load_file as load_safetensors_file
from huggingface_hub import hf_hub_download
-from ..optimization.memory_manager import preinitialize_rope_cache
-from ..common.config import load_config, create_object
-from ..core.infer import VideoDiffusionInfer
+from modules.seedvr.src.optimization.memory_manager import preinitialize_rope_cache
+from modules.seedvr.src.common.config import load_config, create_object
+from modules.seedvr.src.core.infer import VideoDiffusionInfer
def configure_runner(model_name, cache_dir, device:str='cpu', dtype:torch.dtype=None):
diff --git a/modules/seedvr/src/models/dit/nadit.py b/modules/seedvr/src/models/dit/nadit.py
index a67695206..883ca3b56 100644
--- a/modules/seedvr/src/models/dit/nadit.py
+++ b/modules/seedvr/src/models/dit/nadit.py
@@ -68,7 +68,7 @@ class NaDiT(nn.Module):
mlp_type: str = "normal",
window: Optional[Tuple] = None,
window_method: Optional[Tuple[str]] = None,
- temporal_window_size: int = None,
+ temporal_window_size: int | None = None,
temporal_shifted: bool = False,
**kwargs,
):
@@ -172,7 +172,7 @@ class NaDiT(nn.Module):
# Body
cache = Cache(disable=disable_cache)
- for i, block in enumerate(self.blocks):
+ for _i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
@@ -220,7 +220,7 @@ class NaDiTUpscaler(nn.Module):
mlp_type: str = "normal",
window: Optional[Tuple] = None,
window_method: Optional[Tuple[str]] = None,
- temporal_window_size: int = None,
+ temporal_window_size: int | None = None,
temporal_shifted: bool = False,
**kwargs,
):
@@ -334,7 +334,7 @@ class NaDiTUpscaler(nn.Module):
# Body
cache = Cache(disable=disable_cache)
- for i, block in enumerate(self.blocks):
+ for _i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
diff --git a/modules/seedvr/src/models/dit_v2/nadit.py b/modules/seedvr/src/models/dit_v2/nadit.py
index 1f37ef8c3..95e907b8f 100644
--- a/modules/seedvr/src/models/dit_v2/nadit.py
+++ b/modules/seedvr/src/models/dit_v2/nadit.py
@@ -216,7 +216,7 @@ class NaDiT(nn.Module):
emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
# Body
- for i, block in enumerate(self.blocks):
+ for _i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
diff --git a/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py b/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
index 2b07fb1d3..26f85d882 100644
--- a/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
+++ b/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
@@ -616,7 +616,7 @@ class UNetMidBlock3D(nn.Module):
self.resnets = nn.ModuleList(resnets)
def forward(self, hidden_states, temb=None, memory_state: MemoryState = MemoryState.DISABLED):
- video_length, frame_height, frame_width = hidden_states.size()[-3:]
+ video_length, _frame_height, _frame_width = hidden_states.size()[-3:]
hidden_states = self.resnets[0](hidden_states, temb, memory_state=memory_state)
for attn, resnet in zip(self.attentions, self.resnets[1:]):
if attn is not None:
@@ -672,7 +672,7 @@ class Encoder3D(nn.Module):
mid_block_add_attention=True,
# [Override] add extra_cond_dim, temporal down num
temporal_down_num: int = 2,
- extra_cond_dim: int = None,
+ extra_cond_dim: int | None = None,
gradient_checkpoint: bool = False,
inflation_mode: _inflation_mode_t = "tail",
time_receptive_field: _receptive_field_t = "half",
diff --git a/modules/seedvr/src/optimization/memory_manager.py b/modules/seedvr/src/optimization/memory_manager.py
index 39feada97..578b6c340 100644
--- a/modules/seedvr/src/optimization/memory_manager.py
+++ b/modules/seedvr/src/optimization/memory_manager.py
@@ -6,8 +6,8 @@ Extracted from: seedvr2.py (lines 373-405, 607-626, 1016-1044)
"""
import torch
-from ..common.cache import Cache
-from ..models.dit_v2.rope import RotaryEmbeddingBase
+from modules.seedvr.src.common.cache import Cache
+from modules.seedvr.src.models.dit_v2.rope import RotaryEmbeddingBase
def preinitialize_rope_cache(runner) -> None:
@@ -86,10 +86,10 @@ def clear_rope_cache(runner) -> None:
"""
if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'):
# Count entries before cleanup
- cache_size = len(runner.cache.cache)
+ len(runner.cache.cache)
# Free all tensors from cache
- for key, value in runner.cache.cache.items():
+ for _key, value in runner.cache.cache.items():
if isinstance(value, (tuple, list)):
for item in value:
if hasattr(item, 'cpu'):
diff --git a/modules/seedvr/src/utils/color_fix.py b/modules/seedvr/src/utils/color_fix.py
index efe80b67d..7d95cb50d 100644
--- a/modules/seedvr/src/utils/color_fix.py
+++ b/modules/seedvr/src/utils/color_fix.py
@@ -2,7 +2,7 @@ import torch
from PIL import Image
from torch import Tensor
from torch.nn import functional as F
-from ..common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation
+from modules.seedvr.src.common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation
from torchvision.transforms import ToTensor, ToPILImage
def adain_color_fix(target: Image.Image, source: Image.Image):
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 3fd64a2c8..8cdd3d396 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -12,6 +12,7 @@ pipelines = {
'Stable Diffusion Inpaint': getattr(diffusers, 'StableDiffusionInpaintPipeline', None),
'Stable Diffusion Instruct': getattr(diffusers, 'StableDiffusionInstructPix2PixPipeline', None),
'Stable Diffusion 1.5': getattr(diffusers, 'StableDiffusionPipeline', None),
+ 'Stable Diffusion 2': getattr(diffusers, 'StableDiffusionPipeline', None),
'Stable Diffusion 2.x': getattr(diffusers, 'StableDiffusionPipeline', None),
'Stable Diffusion Upscale': getattr(diffusers, 'StableDiffusionUpscalePipeline', None),
'Stable Diffusion XL': getattr(diffusers, 'StableDiffusionXLPipeline', None),
@@ -26,16 +27,17 @@ pipelines = {
'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None),
'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
'FLUX': getattr(diffusers, 'FluxPipeline', None),
+ 'UltraFlux': getattr(diffusers, 'DiffusionPipeline', None),
'FLEX': getattr(diffusers, 'AutoPipelineForText2Image', None),
'Chroma': getattr(diffusers, 'ChromaPipeline', None),
'Sana': getattr(diffusers, 'SanaPipeline', None),
'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
- # 'Lumina-DiMOO': getattr(diffusers, 'LuminaDiMOOPipeline', None),
'Lumina 2': getattr(diffusers, 'Lumina2Pipeline', None),
'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
'Kandinsky 2.1': getattr(diffusers, 'KandinskyCombinedPipeline', None),
'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22CombinedPipeline', None),
'Kandinsky 3.0': getattr(diffusers, 'Kandinsky3Pipeline', None),
+ 'Kandinsky 5.0': getattr(diffusers, 'Kandinsky5T2IPipeline', None),
'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None),
'Kolors': getattr(diffusers, 'KolorsPipeline', None),
'CogView 3': getattr(diffusers, 'CogView3PlusPipeline', None),
@@ -47,6 +49,7 @@ pipelines = {
'Cosmos': getattr(diffusers, 'Cosmos2TextToImagePipeline', None),
'WanAI': getattr(diffusers, 'WanPipeline', None),
'Qwen': getattr(diffusers, 'QwenImagePipeline', None),
+ 'Joy': getattr(diffusers, 'JoyImageEditPipeline', None),
'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None),
'ERNIE-Image': getattr(diffusers, 'ErnieImagePipeline', None),
'Nucleus-Image': getattr(diffusers, 'NucleusMoEImagePipeline', None),
@@ -55,7 +58,6 @@ pipelines = {
'FLUX2 Klein': getattr(diffusers, 'Flux2KleinPipeline', None),
'LongCat': getattr(diffusers, 'LongCatImagePipeline', None),
'GLM-Image': getattr(diffusers, 'GlmImagePipeline', None),
- 'VIBE': getattr(diffusers, 'VIBESanaEditingPipeline', getattr(diffusers, 'DiffusionPipeline', None)),
# dynamically imported and redefined later
'Meissonic': getattr(diffusers, 'DiffusionPipeline', None),
'OmniGen2': getattr(diffusers, 'DiffusionPipeline', None),
@@ -69,6 +71,8 @@ pipelines = {
'Anima': getattr(diffusers, 'DiffusionPipeline', None),
'SDXS': getattr(diffusers, 'DiffusionPipeline', None),
'Step1X-Edit': getattr(diffusers, 'DiffusionPipeline', None),
+ 'VIBE': getattr(diffusers, 'DiffusionPipeline', None),
+ 'Lumina-DiMOO': getattr(diffusers, 'DiffusionPipeline', None),
}
diff --git a/modules/shared_state.py b/modules/shared_state.py
index 4bfe6fae5..41e4bf7f0 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -270,9 +270,8 @@ class State:
def do_set_current_image(self):
from modules import shared, images, sd_samplers_common
- if (self.current_latent is None) or self.disable_preview or (self.preview_job == self.job_no):
+ if self.disable_preview or (self.preview_job == self.job_no):
return False
- self.preview_job = self.job_no
if (shared.opts.show_progress_type == "None") and (shared.history.last_image is not None):
last_image = images.image_grid(shared.history.last_image)
@@ -280,27 +279,37 @@ class State:
self.preview_job = -1
return True
- try:
- sample = self.current_latent
- self.current_image_sampling_step = self.sampling_step
+ if self.current_latent is not None:
try:
- if self.current_noise_pred is not None and self.current_sigma is not None and self.current_sigma_next is not None:
- original_sample = sample - (self.current_noise_pred * (self.current_sigma_next-self.current_sigma))
- if self.prediction_type in {"epsilon", "flow_prediction"}:
- sample = original_sample - (self.current_noise_pred * self.current_sigma)
- elif self.prediction_type == "v_prediction":
- sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type
- except Exception:
- pass # ignore sigma errors
- image = sd_samplers_common.samples_to_image_grid(sample)
- self.assign_current_image(image)
+ self.preview_job = self.job_no
+ sample = self.current_latent
+ self.current_image_sampling_step = self.sampling_step
+ try:
+ if self.current_noise_pred is not None and self.current_sigma is not None and self.current_sigma_next is not None:
+ original_sample = sample - (self.current_noise_pred * (self.current_sigma_next-self.current_sigma))
+ if self.prediction_type in {"epsilon", "flow_prediction"}:
+ sample = original_sample - (self.current_noise_pred * self.current_sigma)
+ elif self.prediction_type == "v_prediction":
+ sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type
+ except Exception:
+ pass # ignore sigma errors
+ image = sd_samplers_common.samples_to_image_grid(sample)
+ self.assign_current_image(image)
+ self.preview_job = -1
+ return True
+ except Exception as e:
+ self.preview_job = -1
+ log.error(f'State image: last={self.id_live_preview} step={self.sampling_step} {e}')
+ display(e, 'State image')
+ return False
+ elif self.current_image is not None:
+ self.preview_job = self.job_no
+ self.assign_current_image(self.current_image)
self.preview_job = -1
return True
- except Exception as e:
- self.preview_job = -1
- log.error(f'State image: last={self.id_live_preview} step={self.sampling_step} {e}')
- display(e, 'State image')
- return False
+ else:
+ pass
+ return False
def assign_current_image(self, image):
self.current_image = image
diff --git a/modules/sharpfin/cms.py b/modules/sharpfin/cms.py
index 18f88bfe3..4327e8646 100644
--- a/modules/sharpfin/cms.py
+++ b/modules/sharpfin/cms.py
@@ -131,7 +131,7 @@ def apply_srgb(
flags=flags
)
else:
- img = cast(Image, profileToProfile(
+ img = cast('Image', profileToProfile(
img,
profile,
_SRGB,
@@ -144,14 +144,14 @@ def apply_srgb(
f"unsupported intent on {path} assuming sRGB: {cms_info}",
path=path,
cms_info=cms_info
- ))
+ ), stacklevel=2)
except PyCMSError as ex:
warn(CMSWarning(
f"{ex} on {path}, assuming sRGB: {cms_info}",
path=path,
cms_info=cms_info,
cause=ex,
- ))
+ ), stacklevel=2)
except Exception as ex:
print(f"{ex} on {path}")
diff --git a/modules/sharpfin/sparse_backend.py b/modules/sharpfin/sparse_backend.py
index 2ea9084ea..9f30f671e 100644
--- a/modules/sharpfin/sparse_backend.py
+++ b/modules/sharpfin/sparse_backend.py
@@ -427,7 +427,7 @@ def _dds_kernel(
b_block_offset = tl.load(block_offsets_t + block_inx)
ptr_B = B + b_block_offset * BLOCK_ELEMENTS
- for sub_block_inx in range(nsub_blocks):
+ for _sub_block_inx in range(nsub_blocks):
a = tl.load(ptr_A)
b = tl.load(ptr_B)
@@ -510,7 +510,7 @@ def triton_dds(
trans_B = not rhs.is_contiguous()
trans_A = (lhs.stride(-2) > 1 and lhs.stride(-1) > 1)
- assert trans_A == False, trans_B == False
+ assert not trans_A, not trans_B
assert lhs.shape[-1] <= rhs.shape[0], "incompatible dimensions"
@@ -583,7 +583,7 @@ def _dds_sbsc_kernel(
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
- for block_slice in range(nsub_blocks):
+ for _block_slice in range(nsub_blocks):
a = tl.load(A_block_ptr, eviction_policy='evict_first', boundary_check=(0,), padding_option='zero')
b = tl.load(B_block_ptr, eviction_policy='evict_last')
@@ -787,7 +787,7 @@ def triton_dds_zerorhs_sbsc(
PAD = math.ceil((kernel_window - 0.5) / k)
- offset, block_height, num_blocks, col_width = block_specs
+ offset, block_height, _num_blocks, col_width = block_specs
assert lhs.ndim == 3
CH = lhs.shape[0]
diff --git a/modules/sharpfin/triton_functional.py b/modules/sharpfin/triton_functional.py
index 98986ff75..4ab17d0ef 100644
--- a/modules/sharpfin/triton_functional.py
+++ b/modules/sharpfin/triton_functional.py
@@ -403,7 +403,7 @@ def downscale_sparse(
BLOCK_SIZE: int = 32,
SPARSE_BLOCK_SIZE: int = 64,
) -> torch.Tensor:
- kernel, window = _get_resize_kernel_triton(resize_kernel)
+ _kernel, window = _get_resize_kernel_triton(resize_kernel)
T_W = target_size[-1]
T_H = target_size[-2]
@@ -447,7 +447,7 @@ def downscale_triton(
resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
do_gamma_handling=True,
) -> torch.Tensor:
- kernel, window = _get_resize_kernel_triton(resize_kernel)
+ _kernel, window = _get_resize_kernel_triton(resize_kernel)
y_s_w = compute_coord_grid(target_size[-1], image.shape[-1], window)
y_s_h = compute_coord_grid(target_size[-2], image.shape[-2], window)
@@ -626,7 +626,7 @@ def downscale_sbsc(
BLOCK_SIZE: int = 32,
SPARSE_BLOCK_SIZE: int = 64,
) -> torch.Tensor:
- kernel, window = _get_resize_kernel_triton(resize_kernel)
+ _kernel, window = _get_resize_kernel_triton(resize_kernel)
T_W = target_size[-1]
T_H = target_size[-2]
@@ -673,7 +673,7 @@ def downscale_sbsc_zerorhs(
BLOCK_SIZE: int = 32,
SPARSE_BLOCK_SIZE: int = 64,
) -> torch.Tensor:
- kernel, window = _get_resize_kernel_triton(resize_kernel)
+ _kernel, window = _get_resize_kernel_triton(resize_kernel)
T_W = target_size[-1]
T_H = target_size[-2]
diff --git a/modules/styles.py b/modules/styles.py
index d263598dd..b25fcd91a 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -446,6 +446,8 @@ class StyleDatabase:
log.error(f'Styles invalid: {p.styles}')
return
for style in p.styles:
+ if style is None or style == '':
+ continue
s = self.find_style(style)
if s == self.no_style:
log.warning(f'Apply style: name="{style}" not found')
diff --git a/modules/teacache/teacache_cogvideox.py b/modules/teacache/teacache_cogvideox.py
index 33ed865d6..a78d666c0 100644
--- a/modules/teacache/teacache_cogvideox.py
+++ b/modules/teacache/teacache_cogvideox.py
@@ -34,7 +34,7 @@ def teacache_cog_forward(
"Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
)
- batch_size, num_frames, channels, height, width = hidden_states.shape
+ batch_size, num_frames, _channels, height, width = hidden_states.shape
# 1. Time embedding
timesteps = timestep
@@ -91,7 +91,7 @@ def teacache_cog_forward(
ori_hidden_states = hidden_states.clone()
ori_encoder_hidden_states = encoder_hidden_states.clone()
# 4. Transformer blocks
- for i, block in enumerate(self.transformer_blocks):
+ for _i, block in enumerate(self.transformer_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
@@ -121,7 +121,7 @@ def teacache_cog_forward(
self.previous_residual_encoder = encoder_hidden_states - ori_encoder_hidden_states
else:
# 4. Transformer blocks
- for i, block in enumerate(self.transformer_blocks):
+ for _i, block in enumerate(self.transformer_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
diff --git a/modules/teacache/teacache_hidream.py b/modules/teacache/teacache_hidream.py
index cb3767cab..c0294eb16 100644
--- a/modules/teacache/teacache_hidream.py
+++ b/modules/teacache/teacache_hidream.py
@@ -1,6 +1,6 @@
from typing import Any, Dict, List, Optional, Tuple
from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.utils import logging, deprecate, USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils import deprecate, USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
import torch
import numpy as np
@@ -130,7 +130,7 @@ def teacache_hidream_forward(
else:
# 2. Blocks
ori_hidden_states = hidden_states.clone()
- for bid, block in enumerate(self.double_stream_blocks):
+ for _bid, block in enumerate(self.double_stream_blocks):
cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
cur_encoder_hidden_states = torch.cat(
[initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1
@@ -166,7 +166,7 @@ def teacache_hidream_forward(
)
hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1)
- for bid, block in enumerate(self.single_stream_blocks):
+ for _bid, block in enumerate(self.single_stream_blocks):
cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1)
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -192,7 +192,7 @@ def teacache_hidream_forward(
hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
self.previous_residual = hidden_states - ori_hidden_states
else:
- for bid, block in enumerate(self.double_stream_blocks):
+ for _bid, block in enumerate(self.double_stream_blocks):
cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
cur_encoder_hidden_states = torch.cat(
[initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1
@@ -228,7 +228,7 @@ def teacache_hidream_forward(
)
hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1)
- for bid, block in enumerate(self.single_stream_blocks):
+ for _bid, block in enumerate(self.single_stream_blocks):
cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1)
if torch.is_grad_enabled() and self.gradient_checkpointing:
diff --git a/modules/teacache/teacache_ltx.py b/modules/teacache/teacache_ltx.py
index 23d4897c0..dd84619d1 100644
--- a/modules/teacache/teacache_ltx.py
+++ b/modules/teacache/teacache_ltx.py
@@ -64,7 +64,7 @@ def teacache_ltx_forward(
inp = self.transformer_blocks[0].norm1(inp)
num_ada_params = self.transformer_blocks[0].scale_shift_table.shape[0]
ada_values = self.transformer_blocks[0].scale_shift_table[None, None] + temb_.reshape(batch_size, temb_.size(1), num_ada_params, -1)
- shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+ shift_msa, scale_msa, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = ada_values.unbind(dim=2)
modulated_inp = inp * (1 + scale_msa) + shift_msa
if self.cnt == 0 or self.cnt == self.num_steps-1:
should_calc = True
diff --git a/modules/teacache/teacache_lumina2.py b/modules/teacache/teacache_lumina2.py
index e5cf8d04d..c33449084 100644
--- a/modules/teacache/teacache_lumina2.py
+++ b/modules/teacache/teacache_lumina2.py
@@ -45,7 +45,7 @@ def teacache_lumina2_forward(
attention_mask_for_main_loop_arg = None
if use_mask:
mask = input_to_main_loop.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
- for i, (enc_len, seq_len_val) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+ for i, (_, seq_len_val) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
mask[i, :seq_len_val] = True
attention_mask_for_main_loop_arg = mask
diff --git a/modules/teacache/teacache_mochi.py b/modules/teacache/teacache_mochi.py
index 2be193a77..e7a12cb94 100644
--- a/modules/teacache/teacache_mochi.py
+++ b/modules/teacache/teacache_mochi.py
@@ -32,7 +32,7 @@ def teacache_mochi_forward(
"Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
)
- batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ batch_size, _num_channels, num_frames, height, width = hidden_states.shape
p = self.config.patch_size
post_patch_height = height // p
@@ -61,7 +61,7 @@ def teacache_mochi_forward(
if self.enable_teacache:
inp = hidden_states.clone()
temb_ = temb.clone()
- modulated_inp, gate_msa, scale_mlp, gate_mlp = self.transformer_blocks[0].norm1(inp, temb_)
+ modulated_inp, _gate_msa, _scale_mlp, _gate_mlp = self.transformer_blocks[0].norm1(inp, temb_)
if self.cnt == 0 or self.cnt == self.num_steps-1:
should_calc = True
self.accumulated_rel_l1_distance = 0
@@ -84,7 +84,7 @@ def teacache_mochi_forward(
hidden_states += self.previous_residual
else:
ori_hidden_states = hidden_states.clone()
- for i, block in enumerate(self.transformer_blocks):
+ for _i, block in enumerate(self.transformer_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
@@ -114,7 +114,7 @@ def teacache_mochi_forward(
hidden_states = self.norm_out(hidden_states, temb)
self.previous_residual = hidden_states - ori_hidden_states
else:
- for i, block in enumerate(self.transformer_blocks):
+ for _i, block in enumerate(self.transformer_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
def custom_forward(*inputs):
diff --git a/modules/ui.py b/modules/ui.py
index 1dc1af44b..af3b37cca 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -79,7 +79,7 @@ def create_ui(startup_timer = None) -> gr.Blocks:
interfaces.clear()
shared.opts.ui_disabled = ui_disabled
if len(ui_disabled) > 0:
- log.warning(f'UI disabled: {ui_disabled}')
+ log.info(f'UI disabled: {ui_disabled}')
if 'txt2img' not in ui_disabled:
with gr.Blocks(analytics_enabled=False) as txt2img_interface:
diff --git a/modules/ui_caption.py b/modules/ui_caption.py
index ceb2c2d31..5d4693e61 100644
--- a/modules/ui_caption.py
+++ b/modules/ui_caption.py
@@ -213,7 +213,7 @@ def create_ui():
clip_model = gr.Dropdown([], value=shared.opts.caption_openclip_model, label='CLiP Model', elem_id='clip_clip_model')
ui_common.create_refresh_button(clip_model, openclip.refresh_clip_models, lambda: {"choices": openclip.refresh_clip_models()}, 'clip_models_refresh')
blip_model = gr.Dropdown(list(openclip.caption_models), value=shared.opts.caption_openclip_blip_model, label='Caption Model', elem_id='btN_clip_blip_model')
- clip_mode = gr.Dropdown(openclip.caption_types, label='Mode', value='fast', elem_id='clip_clip_mode')
+ clip_mode = gr.Dropdown(openclip.caption_types, label='CLiP Mode', value='fast', elem_id='clip_clip_mode')
with gr.Accordion(label='Caption: Advanced Options', open=False, visible=True):
with gr.Row():
clip_max_length = gr.Slider(label='clip: max length', value=shared.opts.caption_openclip_max_length, minimum=16, maximum=1024, step=1, elem_id='clip_caption_max_length')
@@ -305,7 +305,7 @@ def create_ui():
movement = gr.Label(elem_id="caption_label_movement", label="Movement", num_top_classes=5, visible=False)
trending = gr.Label(elem_id="caption_label_trending", label="Trending", num_top_classes=5, visible=False)
flavor = gr.Label(elem_id="caption_label_flavor", label="Flavor", num_top_classes=5, visible=False)
- clip_labels_text = gr.Textbox(elem_id="caption_clip_labels_text", label="CLIP Analysis", lines=15, interactive=False, visible=False, show_label=False)
+ clip_labels_text = gr.Textbox(elem_id="caption_clip_labels_text", label="CLiP Analysis", lines=15, interactive=False, visible=False, show_label=False)
with gr.Row(elem_id='copy_buttons_caption'):
copy_caption_buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "extras"])
diff --git a/modules/ui_common.py b/modules/ui_common.py
index b8b03d0a4..a0dcf926b 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -97,14 +97,14 @@ def delete_files(js_data, files, all_files, index):
os.remove(fn)
if fn in all_files:
all_files.remove(fn)
- log.info(f'Delete: image="{fn}"')
+ log.warning(f'Delete: image="{fn}"')
else:
log.warning(f'Delete: image="{fn}" ui mismatch')
base, _ext = os.path.splitext(fn)
desc = f'{base}.txt'
if os.path.exists(desc) and os.path.isfile(desc):
os.remove(desc)
- log.info(f'Delete: text="{fn}"')
+ log.warning(f'Delete: text="{fn}"')
except Exception as e:
log.error(f'Delete: file="{fn}" {e}')
deleted = ', '.join(deleted) if len(deleted) > 0 else 'none'
diff --git a/modules/ui_control.py b/modules/ui_control.py
index da4784403..12a4c9eb0 100644
--- a/modules/ui_control.py
+++ b/modules/ui_control.py
@@ -56,11 +56,11 @@ def return_controls(res, t: float | None = None):
else:
perf = return_stats(t)
if res is None: # no response
- return [None, None, None, None, '', perf]
+ return [None, None, None, '', perf]
elif isinstance(res, str): # error response
- return [None, None, None, None, res, perf]
+ return [None, None, None, res, perf]
elif isinstance(res, tuple): # standard response received as tuple via control_run->yield(output_images, process_image, result_txt)
- preview_image = res[1] # may be None
+ _preview_image = res[1] # may be None
output_image = res[0][0] if isinstance(res[0], list) else res[0] # may be image or list of images
if isinstance(res[0], list):
output_gallery = res[0] if res[0][0] is not None else []
@@ -68,9 +68,9 @@ def return_controls(res, t: float | None = None):
output_gallery = [res[0]] if res[0] is not None else [] # must return list, but can receive single image
result_txt = res[2] if len(res) > 2 else '' # do we have a message
output_video = res[3] if len(res) > 3 else None # do we have a video filename
- return [preview_image, output_image, output_video, output_gallery, result_txt, perf]
+ return [output_image, output_video, output_gallery, result_txt, perf]
else: # unexpected
- return [None, None, None, None, f'Control: Unexpected response: {type(res)}', perf]
+ return [None, None, None, f'Control: Unexpected response: {type(res)}', perf]
def get_units(*values):
@@ -167,12 +167,10 @@ def create_ui(_blocks: gr.Blocks=None):
with gr.Accordion(open=False, label="Input", elem_id="control_input", elem_classes=["small-accordion"]):
with gr.Row():
- show_input = gr.Checkbox(label="Show input", value=True, elem_id="control_show_input")
- show_preview = gr.Checkbox(label="Show preview", value=False, elem_id="control_show_preview")
- with gr.Row():
- input_type = gr.Radio(label="Control input type", choices=['Control only', 'Init image same as control', 'Separate init image'], value='Control only', type='index', elem_id='control_input_type')
+ input_type = gr.Radio(label="Use init image", choices=['No: Control only', '1st: Same as control', '2nd: Separate image'], value='No: Control only', type='index', elem_id='control_input_type')
with gr.Row():
denoising_strength = gr.Slider(minimum=0.00, maximum=0.99, step=0.01, label='Denoising strength', value=0.30, elem_id="control_input_denoising_strength")
+ skip_processing = gr.Checkbox(label="Skip input processing", value=False, elem_id="control_input_skip_processing")
with gr.Accordion(open=False, label="Size", elem_id="control_size", elem_classes=["small-accordion"]):
with gr.Tabs():
@@ -221,7 +219,7 @@ def create_ui(_blocks: gr.Blocks=None):
timer.startup.record('ui-networks')
with gr.Row(elem_id='control-inputs'):
- with gr.Column(scale=9, elem_id='control-input-column', visible=True) as column_input:
+ with gr.Column(scale=9, elem_id='control-input-column', visible=True) as _column_input:
gr.HTML('Input
')
with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-input'):
input_mode = gr.Label(value='select', visible=False)
@@ -241,15 +239,10 @@ def create_ui(_blocks: gr.Blocks=None):
input_folder = gr.File(label="Input", show_label=False, file_count='directory', file_types=['image'], interactive=True, height=gr_height)
with gr.Column(scale=9, elem_id='control-init-column', visible=False) as column_init:
gr.HTML('Init input')
- with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-init'):
- with gr.Tab('Image', id='init-image') as tab_image_init:
- init_image = gr.Image(label="Input", show_label=False, type="pil", interactive=True, tool="editor", height=gr_height, elem_classes=['control-image'])
- with gr.Tab('Video', id='init-video') as tab_video_init:
- init_video = gr.Video(label="Input", show_label=False, interactive=True, height=gr_height, elem_classes=['control-image'])
- with gr.Tab('Batch', id='init-batch') as tab_batch_init:
- init_batch = gr.File(label="Input", show_label=False, file_count='multiple', file_types=['image'], interactive=True, height=gr_height, elem_classes=['control-image'])
- with gr.Tab('Folder', id='init-folder') as tab_folder_init:
- init_folder = gr.File(label="Input", show_label=False, file_count='directory', file_types=['image'], interactive=True, height=gr_height, elem_classes=['control-image'])
+ if (installer.version['kanvas'] == 'disabled') or (installer.version['kanvas'] == 'unavailable'):
+ init_image = gr.Image(label="Input", show_label=False, type="pil", interactive=True, tool="editor", height=gr_height, elem_classes=['control-image'])
+ else:
+ init_image = gr.HTML(value='Kanvas not initialized ', elem_id='kanvas-container')
with gr.Column(scale=9, elem_id='control-output-column', visible=True) as _column_output:
gr.HTML('Output')
with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-output') as output_tabs:
@@ -259,32 +252,19 @@ def create_ui(_blocks: gr.Blocks=None):
output_image = gr.Image(label="Output", show_label=False, type="pil", interactive=False, tool="editor", height=gr_height, elem_id='control_output_image', elem_classes=['control-image'])
with gr.Tab('Video', id='out-video'):
output_video = gr.Video(label="Output", show_label=False, height=gr_height, elem_id='control_output_video', elem_classes=['control-image'])
- with gr.Column(scale=9, elem_id='control-preview-column', visible=False) as column_preview:
- gr.HTML('Preview')
- with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-preview'):
- with gr.Tab('Preview', id='preview-image') as _tab_preview:
- preview_process = gr.Image(label="Preview", show_label=False, type="pil", interactive=False, height=gr_height, visible=True, elem_id='control_preview', elem_classes=['control-image'])
-
from modules.ui_control_elements import create_ui_elements
- create_ui_elements(units, result_txt, preview_process)
+ create_ui_elements(units, result_txt, output_gallery)
with gr.Row(elem_id="control_script_container"):
input_script_args = scripts_manager.scripts_current.setup_ui(parent='control', accordion=True)
- # handlers
- # for btn in input_buttons:
- # btn.click(fn=helpers.copy_input, inputs=[input_mode, btn, input_image, input_resize, input_inpaint], outputs=[input_image, input_resize, input_inpaint], _js='controlInputMode')
- # btn.click(fn=helpers.transfer_input, inputs=[btn], outputs=[input_image, input_resize, input_inpaint] + input_buttons)
-
# hidden button to update gradio control values
for u in units:
controls.extend(u.controls)
btn_update = gr.Button('Update', interactive=True, visible=False, elem_id='control_update')
btn_update.click(fn=get_units, inputs=controls, outputs=[], show_progress='hidden', queue=False)
- show_input.change(fn=lambda x: gr.update(visible=x), inputs=[show_input], outputs=[column_input])
- show_preview.change(fn=lambda x: gr.update(visible=x), inputs=[show_preview], outputs=[column_preview])
input_type.change(fn=lambda x: gr.update(visible=x == 2), inputs=[input_type], outputs=[column_init])
btn_prompt_counter.click(
fn=call_queue.wrap_queued_call(ui_common.update_token_counter),
@@ -299,11 +279,12 @@ def create_ui(_blocks: gr.Blocks=None):
show_progress = 'hidden',
)
+ image_inputs = 5 * [input_image, init_image] # need to repeat controls for kanvas and non-kanvas modes
select_dict = dict(
fn=helpers.select_input,
_js="controlInputMode",
- inputs=[input_mode, input_image, init_image, input_type, input_video, input_batch, input_folder],
- outputs=[output_tabs, preview_process, result_txt, width_before, height_before],
+ inputs=[input_mode, input_type, input_video, input_batch, input_folder] + image_inputs,
+ outputs=[output_tabs, result_txt, width_before, height_before],
show_progress='hidden',
queue=False,
)
@@ -315,7 +296,7 @@ def create_ui(_blocks: gr.Blocks=None):
prompt.submit(**select_dict)
negative.submit(**select_dict)
btn_generate.click(**select_dict)
- for ctrl in [input_image, input_video, input_batch, input_folder, init_image, init_video, init_batch, init_folder, tab_image, tab_video, tab_batch, tab_folder, tab_image_init, tab_video_init, tab_batch_init, tab_folder_init]:
+ for ctrl in [input_image, input_video, input_batch, input_folder, init_image, tab_image, tab_video, tab_batch, tab_folder]:
if hasattr(ctrl, 'change'):
ctrl.change(**select_dict)
if hasattr(ctrl, 'clear'):
@@ -338,14 +319,13 @@ def create_ui(_blocks: gr.Blocks=None):
resize_mode_before, resize_name_before, resize_context_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
resize_mode_after, resize_name_after, resize_context_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
resize_mode_mask, resize_name_mask, resize_context_mask, width_mask, height_mask, scale_by_mask, selected_scale_tab_mask,
- denoising_strength, batch_count, batch_size,
+ denoising_strength, skip_processing, batch_count, batch_size,
enable_hr, hr_sampler_index, hr_denoising_strength, hr_resize_mode, hr_resize_context, hr_upscaler, hr_force, hr_second_pass_steps, hr_scale, hr_resize_x, hr_resize_y, refiner_steps,
refiner_start, refiner_prompt, refiner_negative,
video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
override_script_name, override_script_args, override_settings,
]
output_fields = [
- preview_process,
output_image,
output_video,
output_gallery,
@@ -470,9 +450,9 @@ def create_ui(_blocks: gr.Blocks=None):
generation_parameters_copypaste.register_paste_params_button(bindings)
if (installer.version['kanvas'] == 'disabled') or (installer.version['kanvas'] == 'unavailable'):
- masking.bind_controls([input_image], preview_process, output_image)
+ masking.bind_controls([input_image], output_image)
else:
- masking.bind_kanvas(input_image, preview_process)
+ masking.bind_kanvas(input_image, output_image)
if os.environ.get('SD_CONTROL_DEBUG', None) is not None: # debug only
from modules.control.test import test_processors, test_controlnets, test_adapters, test_xs, test_lite
@@ -484,11 +464,11 @@ def create_ui(_blocks: gr.Blocks=None):
run_test_adapters_btn = gr.Button(value="Test:Adapters", variant='primary', elem_classes=['control-button'])
run_test_lite_btn = gr.Button(value="Test:Control-LLLite", variant='primary', elem_classes=['control-button'])
- run_test_processors_btn.click(fn=test_processors, inputs=[input_image], outputs=[preview_process, output_image, output_video, output_gallery])
- run_test_controlnets_btn.click(fn=test_controlnets, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])
- run_test_xs_btn.click(fn=test_xs, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])
- run_test_adapters_btn.click(fn=test_adapters, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])
- run_test_lite_btn.click(fn=test_lite, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])
+ run_test_processors_btn.click(fn=test_processors, inputs=[input_image], outputs=[output_image, output_video, output_gallery])
+ run_test_controlnets_btn.click(fn=test_controlnets, inputs=[prompt, negative, input_image], outputs=[output_image, output_video, output_gallery])
+ run_test_xs_btn.click(fn=test_xs, inputs=[prompt, negative, input_image], outputs=[output_image, output_video, output_gallery])
+ run_test_adapters_btn.click(fn=test_adapters, inputs=[prompt, negative, input_image], outputs=[output_image, output_video, output_gallery])
+ run_test_lite_btn.click(fn=test_lite, inputs=[prompt, negative, input_image], outputs=[output_image, output_video, output_gallery])
ui_extra_networks.setup_ui(extra_networks_ui, output_gallery)
return [(control_ui, 'Control', 'control')]
diff --git a/modules/ui_control_elements.py b/modules/ui_control_elements.py
index 6c207087a..cc4579940 100644
--- a/modules/ui_control_elements.py
+++ b/modules/ui_control_elements.py
@@ -11,7 +11,7 @@ from modules import shared, ui_components, ui_symbols, ui_common, masking # pyli
from modules import ui_control_helpers as helpers
-def create_ui_elements(units, result_txt, preview_process):
+def create_ui_elements(units, result_txt, output_gallery):
max_units = shared.opts.control_max_units
with gr.Accordion('Control elements', open=False, elem_id="control_elements"):
with gr.Tabs(elem_id='control-tabs') as _tabs_control_type:
@@ -54,7 +54,7 @@ def create_ui_elements(units, result_txt, preview_process):
process_id = process_id,
model_id = model_id,
model_strength = model_strength,
- preview_process = preview_process,
+ preview_process = output_gallery,
preview_btn = preview_btn,
image_upload = image_upload,
image_reuse = image_reuse,
@@ -103,7 +103,7 @@ def create_ui_elements(units, result_txt, preview_process):
process_id = process_id,
model_id = model_id,
model_strength = model_strength,
- preview_process = preview_process,
+ preview_process = output_gallery,
preview_btn = btn_preview,
image_upload = image_upload,
image_reuse = image_reuse,
@@ -150,7 +150,7 @@ def create_ui_elements(units, result_txt, preview_process):
process_id = process_id,
model_id = model_id,
model_strength = model_strength,
- preview_process = preview_process,
+ preview_process = output_gallery,
preview_btn = btn_preview,
image_upload = image_upload,
image_reuse = image_reuse,
@@ -196,7 +196,7 @@ def create_ui_elements(units, result_txt, preview_process):
process_id = process_id,
model_id = model_id,
model_strength = model_strength,
- preview_process = preview_process,
+ preview_process = output_gallery,
preview_btn = btn_preview,
image_upload = image_upload,
image_reuse = image_reuse,
@@ -239,7 +239,7 @@ def create_ui_elements(units, result_txt, preview_process):
process_id = process_id,
model_id = model_id,
model_strength = model_strength,
- preview_process = preview_process,
+ preview_process = output_gallery,
preview_btn = btn_preview,
image_upload = image_upload,
image_reuse = image_reuse,
diff --git a/modules/ui_control_helpers.py b/modules/ui_control_helpers.py
index 36f63d337..45e933f76 100644
--- a/modules/ui_control_helpers.py
+++ b/modules/ui_control_helpers.py
@@ -16,7 +16,6 @@ busy = False # used to synchronize select_input and generate_click
input_source = None
input_init = None
input_mask = None
-input_prev = None
def initialize():
@@ -85,6 +84,8 @@ def get_video(filepath: str):
def process_kanvas(x): # only used when kanvas overrides gr.Image object
image = None
mask = None
+ if x is None:
+ return image, mask
try: # try base64 decode
t0 = time.time()
image_data = x.get('image', '')
@@ -130,33 +131,35 @@ def process_kanvas(x): # only used when kanvas overrides gr.Image object
return image, mask
-def select_input(input_mode, input_image, init_image, init_type, input_video, input_batch, input_folder):
- global busy, input_source, input_init, input_mask, input_prev # pylint: disable=global-statement
+def select_input(input_mode, init_type, input_video, input_batch, input_folder, *args):
+ global busy, input_source, input_init, input_mask # pylint: disable=global-statement
t0 = time.time()
busy = False
- selected_input = input_image # default: Image or Kanvas
+ selected_input = args[0]
+ init_image = args[1]
if input_mode == 'Video':
selected_input = input_video
elif input_mode == 'Batch':
selected_input = input_batch
elif input_mode == 'Folder':
selected_input = input_folder
+ elif input_mode == 'Kanvas':
+ pass # temp assignment only until we process kanvas inputs
+ else:
+ log.error(f'Input: type={input_mode} unrecognized')
+ selected_input = None
size = [gr.update(), gr.update()]
if selected_input is None:
- # log.debug(f'Select input: image={selected_input}')
input_source = None
- return [gr.Tabs.update(), None, ''] + size
- elif selected_input == input_prev:
- # log.debug(f'Select input: image={selected_input} no change')
- return [gr.Tabs.update(), None, ''] + size
- input_prev = selected_input
+ return [gr.Tabs.update(), ''] + size
busy = True
input_type = type(selected_input)
input_mask = None
status = 'Control input | Unknown'
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ res = [gr.Tabs.update(selected='out-gallery'), status]
+
# control inputs
if isinstance(selected_input, Image.Image): # image via upload -> image
if input_mode == 'Outpaint':
@@ -164,33 +167,41 @@ def select_input(input_mode, input_image, init_image, init_type, input_video, in
selected_input, input_mask = masking.outpaint(input_image=selected_input)
input_source = [selected_input]
input_type = 'PIL.Image'
- status = f'Control input | Image | Size {selected_input.width if selected_input else 0}x{selected_input.height if selected_input else 0} | Mode {selected_input.mode if selected_input else "Unknown"}'
+ status = f'Input | Image | Size {selected_input.width if selected_input else 0}x{selected_input.height if selected_input else 0} | Mode {selected_input.mode if selected_input else "Unknown"}'
size = [gr.update(value=selected_input.width), gr.update(value=selected_input.height)]
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ res = [gr.Tabs.update(selected='out-gallery'), status]
elif isinstance(selected_input, dict) and 'kanvas' in selected_input: # kanvas via js -> kanvas dict
- selected_input, input_mask = process_kanvas(selected_input)
- input_source = [selected_input]
+ input_source = []
+ for i, selected in enumerate(args):
+ img, mask = process_kanvas(selected)
+ if img is not None: # use all images
+ input_source.append(img)
+ if i == 0: # use only first mask
+ selected_input = img
+ input_mask = mask
+ if i == 1:
+ init_image = img # control: separate init image
input_type = 'Kanvas'
- status = f'Control input | Kanvas | Size {selected_input.width if selected_input else 0}x{selected_input.height if selected_input else 0} | Mode {selected_input.mode if selected_input else "Unknown"}'
- if selected_input:
- size = [gr.update(value=selected_input.width), gr.update(value=selected_input.height)]
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ status = f'Input | Kanvas | Images {len(input_source)}'
+ if len(input_source) > 0:
+ size = [gr.update(value=input_source[0].width), gr.update(value=input_source[0].height)]
+ res = [gr.Tabs.update(selected='out-gallery'), status]
elif isinstance(selected_input, dict) and 'mask' in selected_input: # inpaint -> dict image+mask
input_mask = selected_input['mask']
selected_input = selected_input['image']
input_source = [selected_input]
input_type = 'PIL.Image'
- status = f'Control input | Image | Size {selected_input.width if selected_input else 0}x{selected_input.height if selected_input else 0} | Mode {selected_input.mode if selected_input else "Unknown"}'
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ status = f'Input | Image | Size {selected_input.width if selected_input else 0}x{selected_input.height if selected_input else 0} | Mode {selected_input.mode if selected_input else "Unknown"}'
+ res = [gr.Tabs.update(selected='out-gallery'), status]
elif isinstance(selected_input, gr.components.image.Image): # not likely
input_source = [selected_input.value]
input_type = 'gr.Image'
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ res = [gr.Tabs.update(selected='out-gallery'), status]
elif isinstance(selected_input, str) and os.path.exists(selected_input): # video via upload > tmp filepath to video
input_source = selected_input
input_type = 'gr.Video'
status = get_video(input_source)
- res = [gr.Tabs.update(selected='out-video'), input_mask, status]
+ res = [gr.Tabs.update(selected='out-video'), status]
elif isinstance(selected_input, list): # batch or folder via upload -> list of tmp filepaths
if hasattr(selected_input[0], 'name'):
input_type = 'tempfiles'
@@ -198,10 +209,11 @@ def select_input(input_mode, input_image, init_image, init_type, input_video, in
else:
input_type = 'files'
input_source = selected_input
- status = f'Control input | Images | Files {len(input_source)}'
- res = [gr.Tabs.update(selected='out-gallery'), input_mask, status]
+ status = f'Input | Images | Files {len(input_source)}'
+ res = [gr.Tabs.update(selected='out-gallery'), status]
else: # unknown
input_source = None
+
if init_type == 0: # Control only
input_init = None
elif init_type == 1: # Init image same as control assigned during runtime
diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py
index e01eb72be..f4a8f9010 100644
--- a/modules/ui_definitions.py
+++ b/modules/ui_definitions.py
@@ -373,7 +373,7 @@ def create_settings(cmd_opts):
"cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler", "Control"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "LLM", "Control", "Upscaler"]}),
"cuda_compile_backend": OptionInfo("inductor" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'migraphx', 'ipex', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai', 'openvino', 'openvino_fx']}),
"cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
- "cuda_compile_options": OptionInfo(["repeated", "fullgraph", "dynamic"] if not cmd_opts.use_openvino else [], "Model compile options", gr.CheckboxGroup, {"choices": ["precompile", "repeated", "fullgraph", "dynamic", "verbose"]}),
+ "cuda_compile_options": OptionInfo(["repeated", "dynamic", "components"] if not cmd_opts.use_openvino else [], "Model compile options", gr.CheckboxGroup, {"choices": ["components", "precompile", "repeated", "fullgraph", "dynamic", "verbose"]}),
"deep_cache_interval": OptionInfo(3, "DeepCache cache interval", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
}))
@@ -684,7 +684,7 @@ def create_settings(cmd_opts):
"autocomplete_active": OptionInfo(False, "Enable Autocomplete", gr.Checkbox, {"visible": False}),
"autocomplete_enabled": OptionInfo([], "Enabled tag autocomplete files", gr.Dropdown, {"multiselect": True, "choices": [], "visible": False}),
"autocomplete_min_chars": OptionInfo(3, "Min autocomplete chars", gr.Slider, {"minimum": 2, "maximum": 6, "step": 1, "visible": False}),
- "autocomplete_replace_underscores": OptionInfo(True, "Replace underscores in autocomplete", gr.Checkbox, {"visible": False}),
+ "autocomplete_keep_underscores": OptionInfo(False, "Keep underscores in autocomplete", gr.Checkbox, {"visible": False}),
"autocomplete_append_comma": OptionInfo(True, "Append comma after autocomplete", gr.Checkbox, {"visible": False}),
# Caption settings (controlled via Caption Tab UI)
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index 6136c915a..3c3bf2817 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -54,15 +54,14 @@ preview_map = None
def init_api():
- def fetch_file(filename: str = ""):
+ def get_thumb(filename: str = ""):
global allowed_dirs # pylint: disable=global-statement
if len(allowed_dirs) == 0:
allowed_dirs = shared.demo.allowed_paths
if filename is None or len(filename) == 0:
return JSONResponse({ "error": "no filename" }, status_code=400)
- if not os.path.exists(filename) or not os.path.isfile(filename):
+ if not os.path.exists(filename) or not os.path.isfile(filename) or os.path.getsize(filename) == 0:
return FileResponse('html/missing.png', headers={"Accept-Ranges": "bytes"})
- # return JSONResponse({ "error": f"file {filename}: not found" }, status_code=404)
if filename.startswith('html/') or filename.startswith('models/'):
return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
if not any(Path(folder).absolute() in Path(filename).absolute().parents for folder in allowed_dirs):
@@ -118,7 +117,7 @@ def init_api():
return JSONResponse(obj)
shared.api.add_api_route("/sdapi/v1/network", get_network, methods=["GET"])
- shared.api.add_api_route("/sdapi/v1/network/thumb", fetch_file, methods=["GET"], auth=False)
+ shared.api.add_api_route("/sdapi/v1/network/thumb", get_thumb, methods=["GET"], auth=False)
shared.api.add_api_route("/sdapi/v1/network/metadata", get_metadata, methods=["GET"])
shared.api.add_api_route("/sdapi/v1/network/info", get_info, methods=["GET"])
shared.api.add_api_route("/sdapi/v1/network/desc", get_desc, methods=["GET"])
@@ -199,26 +198,37 @@ class ExtraNetworksPage:
def link_preview(self, filename: str):
quoted_filename = urllib.parse.quote(filename.replace('\\', '/'))
- mtime = os.path.getmtime(filename) if os.path.exists(filename) else 0
- preview = f"{shared.opts.subpath}/sdapi/v1/network/thumb?filename={quoted_filename}&mtime={mtime}"
+ # mtime = os.path.getmtime(filename) if os.path.exists(filename) else 0
+ # preview = f"{shared.opts.subpath}/sdapi/v1/network/thumb?filename={quoted_filename}&mtime={mtime}"
+ preview = f"{shared.opts.subpath}/sdapi/v1/network/thumb?filename={quoted_filename}"
return preview
def get_exif(self, image: Image.Image):
import piexif
import piexif.helper
+ parameters = ''
try:
- exifinfo = image.getexif()
- if exifinfo is not None and len(exifinfo) > 0:
- return piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
+ info = image.info or {}
+ for key in ('parameters', 'UserComment'):
+ value = info.get(key)
+ if value and str(value).strip():
+ parameters = str(value)
+ break
except Exception:
pass
- try:
- exifinfo = image.info.get('parameters', None)
- if exifinfo is not None and len(exifinfo) > 0:
- return piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
- except Exception:
- pass
- return piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump('', encoding="unicode") } })
+ if not parameters:
+ try:
+ exif_bytes = (image.info or {}).get('exif')
+ if exif_bytes:
+ parsed = piexif.load(exif_bytes)
+ raw = parsed.get('Exif', {}).get(piexif.ExifIFD.UserComment)
+ if raw:
+ decoded = piexif.helper.UserComment.load(raw)
+ if decoded and decoded.strip():
+ parameters = decoded
+ except Exception:
+ pass
+ return piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters, encoding="unicode") } })
def create_thumb(self):
debug(f'EN create-thumb: {self.name}')
@@ -913,7 +923,10 @@ def create_ui(container, button_parent: gr.Button, tabname: str, skip_indexing =
b64str = ui.last_item.preview.split(',',1)[1]
img = Image.open(io.BytesIO(base64.b64decode(b64str)))
elif hasattr(item, 'local_preview') and os.path.exists(item.local_preview):
- img = item.local_preview
+ if os.path.getsize(item.local_preview) < 1024: # sanity check
+ img = page.find_preview_file(item.filename)
+ else:
+ img = item.local_preview
else:
img = page.find_preview_file(item.filename)
diff --git a/modules/ui_gallery.py b/modules/ui_gallery.py
index 9791598b3..c953cb05a 100644
--- a/modules/ui_gallery.py
+++ b/modules/ui_gallery.py
@@ -59,6 +59,7 @@ def create_ui():
with gr.Blocks() as tab:
with gr.Row(elem_id='tab-gallery-sort-buttons'):
sort_buttons = []
+ """
sort_buttons.append(sort_nameA := ToolButton(value=ui_symbols.sort_alpha_asc, elem_classes=['gallery-sort']))
sort_buttons.append(sort_nameD := ToolButton(value=ui_symbols.sort_alpha_dsc, elem_classes=['gallery-sort']))
sort_buttons.append(sort_sizeA := ToolButton(value=ui_symbols.sort_size_asc, elem_classes=['gallery-sort']))
@@ -67,6 +68,11 @@ def create_ui():
sort_buttons.append(sort_resD := ToolButton(value=ui_symbols.sort_num_dsc, elem_classes=['gallery-sort']))
sort_buttons.append(sort_modA := ToolButton(value=ui_symbols.sort_time_asc, elem_classes=['gallery-sort']))
sort_buttons.append(sort_modD := ToolButton(value=ui_symbols.sort_time_dsc, elem_classes=['gallery-sort']))
+ """
+ sort_buttons.append(sort_name := ToolButton(value=ui_symbols.sort_alpha, elem_classes=['gallery-sort']))
+ sort_buttons.append(sort_size := ToolButton(value=ui_symbols.sort_size, elem_classes=['gallery-sort']))
+ sort_buttons.append(sort_res := ToolButton(value=ui_symbols.sort_num, elem_classes=['gallery-sort']))
+ sort_buttons.append(sort_time := ToolButton(value=ui_symbols.sort_time, elem_classes=['gallery-sort']))
gr.Textbox(show_label=False, placeholder='Search', elem_id='tab-gallery-search')
gr.HTML('', elem_id='tab-gallery-status')
gr.HTML('', elem_id='tab-gallery-progress')
@@ -81,6 +87,7 @@ def create_ui():
gallery_images, gen_info, html_info, _html_info_formatted, html_log = ui_common.create_output_panel("gallery")
btn_gallery_image.click(fn=read_media, _js='gallerySendImage', inputs=[html_info], outputs=[gallery_images, gallery_video, html_info, gen_info, html_log])
+ """
sort_nameA.click(fn=None, _js='() => gallerySort("nameA")')
sort_nameD.click(fn=None, _js='() => gallerySort("nameD")')
sort_sizeA.click(fn=None, _js='() => gallerySort("sizeA")')
@@ -89,5 +96,10 @@ def create_ui():
sort_resD.click(fn=None, _js='() => gallerySort("resD")')
sort_modA.click(fn=None, _js='() => gallerySort("modA")')
sort_modD.click(fn=None, _js='() => gallerySort("modD")')
+ """
+ sort_name.click(fn=None, _js='() => gallerySort("name")')
+ sort_size.click(fn=None, _js='() => gallerySort("size")')
+ sort_res.click(fn=None, _js='() => gallerySort("res")')
+ sort_time.click(fn=None, _js='() => gallerySort("mod")')
return [(tab, 'Gallery', 'tab-gallery')]
diff --git a/modules/ui_models.py b/modules/ui_models.py
index c8a981261..57f621154 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -1,5 +1,6 @@
import os
import inspect
+from html import escape
from typing import cast
import gradio as gr
from modules import errors, sd_models, sd_vae, extras, sd_samplers, ui_symbols, modelstats
@@ -12,6 +13,16 @@ from modules.shared import opts, log
extra_ui = []
+def get_folder_size(folder):
+ total_size = 0
+ for dirpath, _dirnames, filenames in os.walk(folder, followlinks=False):
+ for f in filenames:
+ fp = os.path.join(dirpath, f)
+ if not os.path.islink(fp) and os.path.isfile(fp):
+ total_size += os.path.getsize(fp)
+ return round(total_size / 1024 / 1024 / 1024, 3)
+
+
def update_model_hashes():
from modules import sd_unet, sd_checkpoint
unets = {}
@@ -21,6 +32,70 @@ def update_model_hashes():
yield from sd_models.update_model_hashes(model_type='checkpoint')
+def create_models_table(rows: list = []):
+ from modules import sd_detect
+ rows = sorted(rows, key=lambda row: str(getattr(row, 'model_name', '')).lower())
+ html = """
+
+
+
+ Name
+ Family
+ Type
+ Pipeline
+ Size
+ MTime
+ Hash
+
+
+
+
+ {tbody}
+
+
+ """
+ tbody = ''
+ for row in rows:
+ try:
+ f = row.filename
+ stat_size, stat_mtime = modelstats.stat(f)
+ if os.path.isfile(f):
+ typ = os.path.splitext(f)[1][1:]
+ size = round(stat_size / 1024 / 1024 / 1024, 3)
+ elif os.path.isdir(f):
+ typ = 'diffusers'
+ size = get_folder_size(f)
+ else:
+ typ = 'unknown'
+ size = 0
+ guess = 'Stable Diffusion' # set default guess
+ guess = sd_detect.guess_by_size(f, guess)
+ guess = sd_detect.guess_by_name(f, guess)
+ guess, pipeline = sd_detect.guess_by_diffusers(f, guess)
+ guess = sd_detect.guess_variant(f, guess)
+ pipeline = sd_detect.shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
+ model_name = escape(str(row.model_name))
+ pipeline_name = escape(pipeline.__name__ if pipeline else '(unknown)')
+ typ_name = escape(str(typ))
+ guess_name = escape(str(guess))
+ hash_name = escape(str(row.shorthash))
+ mtime_sort = stat_mtime.timestamp() if hasattr(stat_mtime, 'timestamp') else 0
+ tbody += f"""
+
+ {model_name}
+ {typ_name}
+ {guess_name}
+ {pipeline_name}
+ {size:.3f} GB
+ {stat_mtime}
+ {hash_name}
+ \uf530
+
+ """
+ except Exception as e:
+ log.error(f'Model list: row={vars(row)} {e}')
+ return html.format(tbody=tbody)
+
def create_ui():
log.debug('UI initialize: tab=models')
dummy_component = gr.Label(visible=False)
@@ -31,7 +106,7 @@ def create_ui():
with gr.Column(elem_id='models_input_container', scale=3):
- with gr.Tab(label="Current", elem_id="models_current_tab"):
+ with gr.Tab(label="Active Model", elem_id="models_current_tab"):
def create_modules_table(rows: list):
html = """
@@ -96,65 +171,28 @@ def create_ui():
model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_meta])
- with gr.Tab(label="List", elem_id="models_list_tab"):
- def create_models_table(rows: list):
- from modules import sd_detect
- html = """
-
-
- Name Type Detect Pipeline Hash Size MTime
-
-
- {tbody}
-
-
- """
- tbody = ''
- for row in rows:
- try:
- f = row.filename
- stat_size, stat_mtime = modelstats.stat(f)
- if os.path.isfile(f):
- typ = os.path.splitext(f)[1][1:]
- size = f"{round(stat_size / 1024 / 1024 / 1024, 3)} gb"
- elif os.path.isdir(f):
- typ = 'diffusers'
- size = 'folder'
- else:
- typ = 'unknown'
- size = 'unknown'
- guess = 'Diffusion' # set default guess
- guess = sd_detect.guess_by_size(f, guess)
- guess = sd_detect.guess_by_name(f, guess)
- guess, pipeline = sd_detect.guess_by_diffusers(f, guess)
- guess = sd_detect.guess_variant(f, guess)
- pipeline = sd_detect.shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
- tbody += f"""
-
- {row.model_name}
- {typ}
- {guess}
- {pipeline.__name__ if pipeline else '(unknown)'}
- {row.shorthash}
- {size}
- {stat_mtime}
-
- """
- except Exception as e:
- log.error(f'Model list: row={vars(row)} {e}')
- return html.format(tbody=tbody)
-
+ with gr.Tab(label="Models List", elem_id="models_list_tab"):
with gr.Row():
gr.HTML('List all locally available models ')
with gr.Row():
- model_list_btn = gr.Button(value="List models", variant='primary')
- model_checkhash_btn = gr.Button(value="Calculate missing hashes", variant='secondary')
+ model_list_btn = gr.Button(value="Refresh list", variant='primary')
+ model_checkhash_btn = gr.Button(value="Calculate hashes", variant='secondary')
with gr.Row():
- model_table = gr.HTML(value='', elem_id="model_list_table")
+ model_table = gr.HTML(value=create_models_table(), elem_id="model_list_table", elem_classes="scroll-auto")
model_checkhash_btn.click(fn=update_model_hashes, inputs=[], outputs=[model_table])
model_list_btn.click(fn=lambda: create_models_table(list(sd_models.checkpoints_list.values())), inputs=[], outputs=[model_table])
+ with gr.Tab(label="Cache List", elem_id="models_cache_tab"):
+ with gr.Row():
+ gr.HTML('List models in Hugging Face cache ')
+ with gr.Row():
+ model_cache_btn = gr.Button(value="Refresh list", variant='primary')
+ with gr.Row():
+ cache_table = gr.HTML(value=create_models_table(), elem_id="model_cache_table", elem_classes="scroll-auto")
+
+ model_cache_btn.click(fn=lambda: create_models_table(sd_models.list_hfcache()), inputs=[], outputs=[cache_table])
+
with gr.Tab(label="Metadata", elem_id="models_metadata_tab"):
from modules.civitai.metadata_civitai import civit_search_metadata, civit_update_metadata
with gr.Row():
@@ -181,7 +219,7 @@ def create_ui():
return ['None'] + sd_models.checkpoint_titles()
with gr.Row():
- gr.HTML(' Merge multiple models ')
+ gr.HTML(' Merge multiple models ')
with gr.Row(equal_height=False):
with gr.Column(variant='compact'):
with gr.Row():
@@ -422,7 +460,7 @@ def create_ui():
with gr.Tab(label="Replace", elem_id="models_replace_tab"):
with gr.Row():
- gr.HTML(' Replace model components ')
+ gr.HTML(' Replace model components ')
with gr.Row():
with gr.Column(scale=3):
model_type = gr.Dropdown(label="Base model type", choices=['sd15', 'sdxl', 'sd21', 'sd35', 'flux.1'], value='sdxl', interactive=False)
@@ -628,7 +666,7 @@ def create_ui():
from modules.models_hf import hf_search, hf_select, hf_download_model, hf_update_token
with gr.Column(scale=6):
with gr.Row():
- gr.HTML(' Download model from huggingface ')
+ gr.HTML(' Download model from huggingface ')
with gr.Row():
hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models')
hf_search_btn = ToolButton(value=ui_symbols.search, interactive=True, elem_id="hf_text_search")
diff --git a/modules/ui_models_load.py b/modules/ui_models_load.py
index 9e6ca45cf..9246a10d6 100644
--- a/modules/ui_models_load.py
+++ b/modules/ui_models_load.py
@@ -276,7 +276,7 @@ def create_ui(gr_status, gr_file):
return 'Save receipe not implemented yet'
with gr.Row():
- gr.HTML('')
+ gr.HTML(' Custom model loader ')
with gr.Row():
choices = list(shared_items.pipelines)
choices = ['Current' if x.startswith('Custom') else x for x in choices]
diff --git a/modules/ui_settings.py b/modules/ui_settings.py
index f6401a9fe..cbb26c1b5 100644
--- a/modules/ui_settings.py
+++ b/modules/ui_settings.py
@@ -140,11 +140,14 @@ def run_settings(*args):
directml_override_opts()
if shared.cmd_opts.use_openvino:
if "Model" not in shared.opts.cuda_compile:
- log.warning("OpenVINO: Enabling Torch Compile Model")
+ log.warning("OpenVINO: Overriding Torch Compile Model")
shared.opts.cuda_compile.append("Model")
if shared.opts.cuda_compile_backend != shared.opts.openvino_compile_backend:
- log.warning(f"OpenVINO: Enabling Torch Compile backend={shared.opts.openvino_compile_backend}")
+ log.warning(f"OpenVINO: Overriding Torch Compile backend={shared.opts.openvino_compile_backend}")
shared.opts.cuda_compile_backend = shared.opts.openvino_compile_backend
+ if shared.opts.diffusers_offload_mode != "none":
+ log.warning("OpenVINO: Overriding diffusers_offload_mode=none")
+ shared.opts.diffusers_offload_mode = "none"
if shared.opts.sd_backend != "diffusers":
log.error('Legacy option: backend=original is no longer supported')
shared.opts.sd_backend = "diffusers"
@@ -157,10 +160,10 @@ def run_settings(*args):
return shared.opts.dumpjson(), f'{len(changed)} Settings changed without save: {", ".join(changed)}'
return shared.opts.dumpjson(), f'{len(changed)} Settings changed{": " if len(changed) > 0 else ""}{", ".join(changed)}'
-def run_settings_single(value, key, progress=False):
+def run_settings_single(value, key, progress=False, force=False):
if not shared.opts.same_type(value, shared.opts.data_labels[key].default):
return gr.update(visible=True), shared.opts.dumpjson()
- if not shared.opts.set(key, value):
+ if not shared.opts.set(key, value, force):
return gr.update(value=getattr(shared.opts, key)), shared.opts.dumpjson()
if key == "cuda_compile_backend" and value == "olive-ai":
from modules.onnx_impl import install_olive
@@ -168,9 +171,9 @@ def run_settings_single(value, key, progress=False):
if shared.cmd_opts.use_directml:
from modules.dml import directml_override_opts
directml_override_opts()
- shared.opts.save()
- if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet']:
- log.debug(f'Setting changed: {key}={value} progress={progress}')
+ shared.opts.save(silent=True)
+ if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet'] or force:
+ log.debug(f'Setting changed: {key}="{value}" progress={progress} force={force}')
return get_value_for_setting(key), shared.opts.dumpjson()
@@ -203,8 +206,8 @@ def create_ui(disabled_tabs=None):
unload_sd_model.click(fn=unload_sd_weights, inputs=[], outputs=[])
reload_sd_model.click(fn=reload_sd_weights, inputs=[], outputs=[])
enable_profiling.click(fn=switch_profiling, inputs=[], outputs=[enable_profiling])
- restart_submit.click(fn=lambda: shared.restart_server(restart=True), _js="restartReload")
- shutdown_submit.click(fn=lambda: shared.restart_server(restart=False), _js="restartReload")
+ restart_submit.click(fn=lambda: shared.restart_server(restart=True), _js="() => { restartReload(); }")
+ shutdown_submit.click(fn=lambda: shared.restart_server(restart=False), _js="() => { restartReload(); }")
with gr.Tabs(elem_id="system") as system_tabs:
global ui_system_tabs # pylint: disable=global-statement
@@ -310,7 +313,7 @@ def create_ui(disabled_tabs=None):
outputs=[text_settings, result],
)
if defaults_submit:
- defaults_submit.click(fn=lambda: shared.restore_defaults(restart=True), _js="restartReload")
+ defaults_submit.click(fn=lambda: shared.restore_defaults(restart=True), _js="() => { restartReload(); }")
def reset_quicksettings(quick_components):
@@ -377,7 +380,7 @@ def create_quicksettings(interfaces):
button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False)
button_set_checkpoint.click(
- fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'),
+ fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint', force=True),
_js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }",
inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component],
outputs=[shared.settings_components['sd_model_checkpoint'], text_settings],
diff --git a/modules/ui_symbols.py b/modules/ui_symbols.py
index 6819cb5f6..1a82cfc31 100644
--- a/modules/ui_symbols.py
+++ b/modules/ui_symbols.py
@@ -38,12 +38,16 @@ bullet = '⃝'
vision = '\uf06e' # Font Awesome eye icon (more minimalistic)
reasoning = '\uf0eb' # Font Awesome lightbulb icon (represents thinking/reasoning)
cloud = '☁'
+sort_alpha = '\uf15d'
sort_alpha_asc = '\uf15d'
sort_alpha_dsc = '\uf15e'
+sort_size = '\uf160'
sort_size_asc = '\uf160'
sort_size_dsc = '\uf161'
+sort_num = '\uf162'
sort_num_asc = '\uf162'
sort_num_dsc = '\uf163'
+sort_time = '\uf0dc'
sort_time_asc = '\uf0de'
sort_time_dsc = '\uf0dd'
style_apply = '↶'
diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py
index 1ee384fdf..5c5c299af 100644
--- a/modules/upscaler_simple.py
+++ b/modules/upscaler_simple.py
@@ -99,5 +99,5 @@ class UpscalerLatent(Upscaler):
elif selected_model == "Latent Bicubic antialias":
mode, antialias = 'bicubic', True
else:
- raise log.error(f"Upscale: type=latent model={selected_model} unknown")
+ log.error(f"Upscale: type=latent model={selected_model} unknown")
return F.interpolate(img, size=(h, w), mode=mode, antialias=antialias)
diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py
index 7e228d9c0..aedbf799e 100644
--- a/modules/vae/sd_vae_taesd.py
+++ b/modules/vae/sd_vae_taesd.py
@@ -45,7 +45,6 @@ prev_cls = ''
prev_type = ''
prev_model = ''
lock = threading.Lock()
-supported = ['sd', 'sdxl', 'sd3', 'f1', 'f2', 'h1', 'zimage', 'lumina2', 'hunyuanvideo', 'wanai', 'chrono', 'cosmos', 'anima', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4']
def warn_once(msg, variant=None):
@@ -57,16 +56,19 @@ def warn_once(msg, variant=None):
return Image.new('RGB', (8, 8), color = (0, 0, 0))
-def get_model(model_type = 'decoder', variant = None):
- global prev_cls, prev_type, prev_model, prev_warnings # pylint: disable=global-statement
- model_cls = shared.sd_model_type
- if model_cls is None or model_cls == 'none':
- return None, variant
- elif model_cls in {'ldm', 'pixartalpha'}:
+def get_model(model_cls, variant=None):
+ if variant is not None:
+ pass
+ if model_cls in {'sd'}:
model_cls = 'sd'
+ variant = shared.opts.taesd_variant
+ elif model_cls in {'sdxl', 'ldm', 'pixartalpha'}:
+ model_cls = 'sdxl'
+ variant = shared.opts.taesd_variant
elif model_cls in {'pixartsigma', 'hunyuandit', 'omnigen', 'auraflow'}:
model_cls = 'sdxl'
- elif model_cls in {'f1', 'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4'}:
+ variant = shared.opts.taesd_variant
+ elif model_cls in {'f1', 'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4', 'ultraflux'}:
model_cls = 'f1'
variant = 'TAE FLUX.1'
elif model_cls in {'f2', 'ernieimage'}:
@@ -74,12 +76,24 @@ def get_model(model_type = 'decoder', variant = None):
variant = 'TAE FLUX.2'
elif model_cls in {'sd3'}:
variant = 'TAE SD3'
- elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos', 'anima', 'fibo'}:
- variant = variant or 'TAE WanVideo'
- elif model_cls not in supported:
+ elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos', 'anima', 'fibo', 'joy'}:
+ variant = 'TAE WanVideo'
+ else:
warn_once(f'cls={shared.sd_model.__class__.__name__} type={shared.sd_model_type} unsuppported', variant=variant)
+ return model_cls, None
+ if debug:
+ log.debug(f'TAESD detect: cls={model_cls} variant={variant}')
+ return model_cls, variant
+
+
+def load_model(model_type = 'decoder', variant = None):
+ global prev_cls, prev_type, prev_model, prev_warnings # pylint: disable=global-statement
+ model_cls = shared.sd_model_type
+ if model_cls is None or model_cls == 'none':
+ return None, variant
+ model_cls, variant = get_model(model_cls, variant)
+ if model_cls is None or variant is None:
return None, variant
- variant = variant or shared.opts.taesd_variant
folder = os.path.join(paths.models_path, "TAESD")
dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16
os.makedirs(folder, exist_ok=True)
@@ -154,7 +168,7 @@ def decode(latents):
global first_run # pylint: disable=global-statement
with lock:
try:
- vae, variant = get_model(model_type='decoder')
+ vae, variant = load_model(model_type='decoder')
if vae is None or max(latents.shape) > 256: # safetey check of large tensors
return latents
except Exception as e:
@@ -192,7 +206,7 @@ def decode(latents):
def encode(image):
with lock:
- vae, variant = get_model(model_type='encoder')
+ vae, variant = load_model(model_type='encoder')
if vae is None:
return image
try:
diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py
index db01c605b..f63ae5d3b 100644
--- a/modules/video_models/video_save.py
+++ b/modules/video_models/video_save.py
@@ -80,22 +80,25 @@ def numpy_to_tensor(images):
return tensor
+def add_audio_stream(container, audio_sample_rate: int):
+ # Must be registered before the first container.mux(); avformat_write_header runs there
+ # and freezes the stream set, after which new streams have time_base=0/0.
+ audio_stream = container.add_stream("aac", rate=audio_sample_rate)
+ audio_stream.codec_context.sample_rate = audio_sample_rate
+ audio_stream.codec_context.layout = "stereo"
+ audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
+ log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} base={audio_stream.codec_context.time_base}')
+ return audio_stream
+
+
def write_audio(
container,
+ audio_stream,
samples: torch.Tensor,
audio_sample_rate: int,
) -> None:
av = check_av()
- # create stream
- audio_options = { 'time_base': f'1/{audio_sample_rate}' }
- audio_stream = container.add_stream("aac", rate=audio_sample_rate, options=audio_options)
- audio_stream.codec_context.sample_rate = audio_sample_rate
- audio_stream.codec_context.layout = "stereo"
audio_stream.codec_context.format = "fltp"
- audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
- # audio_stream.time_base = audio_stream.codec_context.time_base # TODO audio set time-base
- log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} format={audio_stream.codec_context.format} base={audio_stream.codec_context.time_base}')
- # init input samples
if samples.ndim == 1:
samples = samples[:, None]
if samples.shape[1] != 2 and samples.shape[0] == 2:
@@ -111,13 +114,11 @@ def write_audio(
layout="stereo",
)
audio_frames.sample_rate = audio_sample_rate
- # init resampler
audio_resampler = av.audio.resampler.AudioResampler(
format=audio_stream.codec_context.format,
layout=audio_stream.codec_context.layout,
rate=audio_stream.codec_context.sample_rate,
)
- # resample
pts = 0
for resampled in audio_resampler.resample(audio_frames):
resampled.pts = resampled.pts or 0
@@ -126,7 +127,6 @@ def write_audio(
for packet in packets:
container.mux(packet)
pts += resampled.samples
- # flush audio encoder
for packet in audio_stream.encode():
container.mux(packet)
@@ -177,6 +177,7 @@ def atomic_save_video(
stream.width = video_array.shape[2]
stream.height = video_array.shape[1]
stream.pix_fmt = pix_fmt
+ audio_stream = add_audio_stream(container, aac) if audio is not None else None
for img in video_array:
frame = av.VideoFrame.from_ndarray(img, format="rgb24")
for packet in stream.encode_lazy(frame):
@@ -185,9 +186,9 @@ def atomic_save_video(
pbar.update(task, advance=1)
for packet in stream.encode(): # flush
container.mux(packet)
- if audio is not None:
+ if audio_stream is not None:
try:
- write_audio(container, audio, aac)
+ write_audio(container, audio_stream, audio, aac)
except Exception as e:
log.error(f'Video audio encoding: {e}')
errors.display(e, 'Audio')
diff --git a/package.json b/package.json
index 0bfe694d5..b280e6231 100644
--- a/package.json
+++ b/package.json
@@ -30,8 +30,8 @@
"pylint-win": "venv\\scripts\\activate && pylint *.py modules/ pipelines/ scripts/ extensions-builtin/",
"pyright": ". venv/bin/activate && pyright --threads 4",
"pyright-win": "venv\\scripts\\activate && pyright --threads 4",
- "ty": ". venv/bin/activate && ty check",
- "ty-win": "venv\\scripts\\activate && ty check",
+ "ty": ". venv/bin/activate && ty check --force-exclude",
+ "ty-win": "venv\\scripts\\activate && ty check --force-exclude",
"lint": "npm run format && npm run eslint && npm run eslint-ui && npm run eslint-kanvas && npm run ruff && npm run pylint",
"lint-win": "npm run format-win && npm run eslint && npm run eslint-ui && npm run eslint-kanvas && npm run ruff-win && npm run pylint-win",
"test": ". venv/bin/activate; python launch.py --debug --test",
diff --git a/pipelines/anima/anima_image.py b/pipelines/anima/anima_image.py
new file mode 100644
index 000000000..5d0334181
--- /dev/null
+++ b/pipelines/anima/anima_image.py
@@ -0,0 +1,168 @@
+"""Anima img2img and inpainting pipelines (built dynamically from the runtime-imported base class)."""
+
+from typing import Callable, Dict, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils.torch_utils import randn_tensor
+
+from modules import devices
+
+
+def _encode_image(pipe, image, dtype, device, generator):
+ """VAE-encode an image and normalize to denoiser latent space."""
+ if isinstance(image, list):
+ image = image[0]
+ image_tensor = pipe.video_processor.preprocess(image, None, None)
+ image_tensor = image_tensor.squeeze(0).to(device=device, dtype=pipe.vae.dtype)
+ image_tensor = image_tensor[None, :, None, :, :]
+ init_latents = pipe.vae.encode(image_tensor).latent_dist.sample(generator)
+ latents_mean = torch.tensor(pipe.vae.config.latents_mean, device=device, dtype=torch.float32).view(1, pipe.vae.config.z_dim, 1, 1, 1)
+ latents_std_inv = (1.0 / torch.tensor(pipe.vae.config.latents_std, device=device, dtype=torch.float32)).view(1, pipe.vae.config.z_dim, 1, 1, 1)
+ return ((init_latents.float() - latents_mean) * latents_std_inv).to(dtype)
+
+
+def _setup_img2img_schedule(scheduler, strength, num_inference_steps, device):
+ """Set custom sigma schedule, return first sigma after scheduler shift."""
+ custom_sigmas = torch.linspace(max(strength, 0.01), 0.0, num_inference_steps).tolist()
+ scheduler.set_timesteps(sigmas=custom_sigmas, device=device)
+ return scheduler.sigmas[0].item()
+
+
+def build_anima_pipeline_classes(base_cls):
+ """Return (AnimaImageToImagePipeline, AnimaInpaintPipeline) inheriting from base_cls."""
+
+ class AnimaImageToImagePipeline(base_cls):
+ """Anima img2img pipeline."""
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ image: Optional[PipelineImageInput] = None,
+ strength: float = 0.8,
+ height: int = 768,
+ width: int = 1360,
+ num_inference_steps: int = 35,
+ guidance_scale: float = 7.0,
+ num_images_per_prompt: Optional[int] = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ max_sequence_length: int = 512,
+ ):
+ actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+ init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+ noise = randn_tensor(init_latents.shape, generator=generator, device=devices.device, dtype=devices.dtype)
+ noised = (actual_sigma * noise + (1.0 - actual_sigma) * init_latents).to(torch.float32)
+
+ orig_set_timesteps = self.scheduler.set_timesteps
+ self.scheduler.set_timesteps = lambda *args, **kwargs: None
+ try:
+ return super().__call__(
+ prompt=prompt, negative_prompt=negative_prompt, height=height, width=width,
+ num_inference_steps=num_inference_steps, guidance_scale=guidance_scale,
+ num_images_per_prompt=num_images_per_prompt, generator=generator, latents=noised,
+ prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds,
+ output_type=output_type, return_dict=return_dict,
+ callback_on_step_end=callback_on_step_end, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ max_sequence_length=max_sequence_length,
+ )
+ finally:
+ self.scheduler.set_timesteps = orig_set_timesteps
+
+ class AnimaInpaintPipeline(AnimaImageToImagePipeline):
+ """Anima inpainting pipeline."""
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ image: Optional[PipelineImageInput] = None,
+ mask_image: Optional[PipelineImageInput] = None,
+ strength: float = 0.8,
+ height: int = 768,
+ width: int = 1360,
+ num_inference_steps: int = 35,
+ guidance_scale: float = 7.0,
+ num_images_per_prompt: Optional[int] = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ max_sequence_length: int = 512,
+ ):
+ actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+ init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+ noise = randn_tensor(init_latents.shape, generator=generator, device=devices.device, dtype=devices.dtype)
+ noised = (actual_sigma * noise + (1.0 - actual_sigma) * init_latents).to(torch.float32)
+ mask_latent = _prepare_mask(self, mask_image, height, width, devices.device)
+
+ orig_set_timesteps = self.scheduler.set_timesteps
+ self.scheduler.set_timesteps = lambda *args, **kwargs: None
+
+ user_callback = callback_on_step_end
+
+ def blend_callback(pipe, i, t, callback_kwargs):
+ cur_latents = callback_kwargs.get("latents")
+ if cur_latents is not None:
+ sigma_next = pipe.scheduler.sigmas[i + 1].item() if i + 1 < len(pipe.scheduler.sigmas) else 0.0
+ init_at_t = sigma_next * noise + (1.0 - sigma_next) * init_latents
+ blended = mask_latent * cur_latents + (1.0 - mask_latent) * init_at_t.to(cur_latents.dtype)
+ callback_kwargs["latents"] = blended
+ if user_callback is not None:
+ callback_kwargs = user_callback(pipe, i, t, callback_kwargs)
+ return callback_kwargs
+
+ try:
+ return base_cls.__call__(
+ self,
+ prompt=prompt, negative_prompt=negative_prompt,
+ height=height, width=width, num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale, num_images_per_prompt=num_images_per_prompt,
+ generator=generator, latents=noised, prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds, output_type=output_type,
+ return_dict=return_dict, callback_on_step_end=blend_callback,
+ callback_on_step_end_tensor_inputs=["latents"],
+ max_sequence_length=max_sequence_length,
+ )
+ finally:
+ self.scheduler.set_timesteps = orig_set_timesteps
+
+ return AnimaImageToImagePipeline, AnimaInpaintPipeline
+
+
+def _prepare_mask(pipe, mask_image, height, width, device):
+ if isinstance(mask_image, Image.Image):
+ mask_image = mask_image.convert("L")
+ if isinstance(mask_image, Image.Image):
+ import torchvision.transforms.functional as TF
+ mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).to(device=device, dtype=torch.float32)
+ elif isinstance(mask_image, torch.Tensor):
+ mask_tensor = mask_image.to(device=device, dtype=torch.float32)
+ if mask_tensor.ndim == 2:
+ mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
+ elif mask_tensor.ndim == 3:
+ mask_tensor = mask_tensor.unsqueeze(0)
+ else:
+ mask_tensor = torch.ones(1, 1, height, width, device=device, dtype=torch.float32)
+ latent_h = height // pipe.vae_scale_factor_spatial
+ latent_w = width // pipe.vae_scale_factor_spatial
+ mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest")
+ mask_latent = mask_latent[:, :1, :, :]
+ mask_latent = mask_latent.unsqueeze(2)
+ return mask_latent
diff --git a/pipelines/anima/anima_lora.py b/pipelines/anima/anima_lora.py
index d75c40249..8454425de 100644
--- a/pipelines/anima/anima_lora.py
+++ b/pipelines/anima/anima_lora.py
@@ -114,7 +114,7 @@ def try_load_lora(name, network_on_disk, lora_scale):
matched += 1
if matched == 0:
return None
- log.debug(f'Network load: type=LoRA name="{name}" native modules={matched} unmatched={unmatched} scale={lora_scale}')
+ log.debug(f'Network load: type=LoRA name="{name}" method=native modules={matched} unmatched={unmatched} scale={lora_scale}')
if unmatched > 0 and l.debug:
log.debug(f'Network load: type=LoRA name="{name}" unmatched_samples={unmatched_samples}')
l.timer.activate += time.time() - t0
diff --git a/pipelines/bria/bria_pipeline.py b/pipelines/bria/bria_pipeline.py
index 6049e3f45..0c4d160d7 100644
--- a/pipelines/bria/bria_pipeline.py
+++ b/pipelines/bria/bria_pipeline.py
@@ -622,7 +622,7 @@ class BriaPipeline(FluxPipeline):
@staticmethod
def _unpack_latents(latents, height, width, vae_scale_factor):
- batch_size, num_patches, channels = latents.shape
+ batch_size, _num_patches, channels = latents.shape
height = height // vae_scale_factor
width = width // vae_scale_factor
diff --git a/models/Reference/iitolstykh--VIBE-Image-Edit.jpg b/pipelines/chroma/__init__.py
similarity index 100%
rename from models/Reference/iitolstykh--VIBE-Image-Edit.jpg
rename to pipelines/chroma/__init__.py
diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py
new file mode 100644
index 000000000..7434421a3
--- /dev/null
+++ b/pipelines/chroma/chroma_lora.py
@@ -0,0 +1,535 @@
+"""Chroma native adapter loader.
+
+Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
+(``lora_force_diffusers`` off and ``chroma`` in ``allow_native``). Reads the
+safetensors directly and writes into sdnext's existing
+``network_layer_mapping``, returning a ``Network`` populated with
+``NetworkModule*`` entries that ``network_activate`` will apply. If the
+setting is on, the diffusers PEFT path handles the file instead.
+
+Entry points, one per family:
+
+- LoRA (+ DoRA) via :func:`try_load_lora`
+- LoKR via :func:`try_load_lokr`
+- LoHA via :func:`try_load_loha` (fused groups skipped, no chunk variant)
+- OFT via :func:`try_load_oft` (fused groups skipped, no chunk variant)
+
+Recognized key prefixes: ``diffusion_model.``, ``transformer.``, ``lora_unet_``.
+Diffusers-PEFT ``lora_A``/``lora_B`` are normalized to ``lora_down``/``lora_up``.
+
+Chroma LoRAs are trained against the Flux block layout regardless of which
+key style they save in:
+
+- ``double_blocks.{i}.{img,txt}_attn.{proj,qkv}``
+- ``double_blocks.{i}.{img,txt}_mlp.{0,2}``
+- ``single_blocks.{i}.{linear1,linear2}``
+
+The diffusers ``ChromaTransformer2DModel`` exposes split-attention modules at:
+
+- ``transformer_blocks.{i}.attn.{to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out}``
+- ``transformer_blocks.{i}.{ff,ff_context}.net.{0.proj,2}``
+- ``single_transformer_blocks.{i}.attn.{to_q,to_k,to_v}``
+- ``single_transformer_blocks.{i}.{proj_mlp,proj_out}``
+
+This loader path-rewrites Flux-layout keys to diffusers names and expands
+fused QKV at load time (LoRA) or apply time (LoKR via
+:class:`NetworkModuleLokrSliceChunk`). For LoHA/OFT the fused groups are
+skipped with a warning.
+
+Chroma's modulation generator is the central ``distilled_guidance_layer``
+approximator (replacing Flux's per-block ``norm1.linear``). The pruned
+AdaLN classes have no ``.linear`` submodule, so any ``_mod_lin`` /
+``_modulation_lin`` keys naturally land in ``unmapped`` and are reported.
+LoRAs that target the approximator itself work without special casing,
+since ``distilled_guidance_layer.<...>`` is a real module path that
+``assign_network_names_to_compvis_modules`` registers.
+"""
+
+import os
+import time
+import torch
+from modules import shared, sd_models
+from modules.logger import log
+from modules.lora import network, network_lora, network_lokr, network_hada, network_oft, lora_convert
+from modules.lora import lora_common as l
+
+
+KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_")
+
+LORA_SUFFIXES = (
+ ".lora_down.weight", ".lora_up.weight",
+ ".lora_A.weight", ".lora_B.weight",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOKR_SUFFIXES = (
+ ".lokr_w1", ".lokr_w2",
+ ".lokr_w1_a", ".lokr_w1_b",
+ ".lokr_w2_a", ".lokr_w2_b",
+ ".lokr_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOHA_SUFFIXES = (
+ ".hada_w1_a", ".hada_w1_b",
+ ".hada_w2_a", ".hada_w2_b",
+ ".hada_t1", ".hada_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+OFT_SUFFIXES = (
+ ".oft_blocks", ".oft_diag",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+
+LORA_MARKERS = (".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight")
+LOKR_MARKERS = (".lokr_w1", ".lokr_w2")
+LOHA_MARKERS = (".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b")
+OFT_MARKERS = (".oft_blocks", ".oft_diag")
+
+SUFFIX_NORMALIZE = {
+ "lora_A.weight": "lora_down.weight",
+ "lora_B.weight": "lora_up.weight",
+}
+
+# Default block counts for Chroma1-HD; overridden at runtime from the live
+# transformer's config when available.
+DEFAULT_NUM_DOUBLE_LAYERS = 19
+DEFAULT_NUM_SINGLE_LAYERS = 38
+
+# Fused QKV split dims. Single blocks fuse Q/K/V plus proj_mlp into linear1;
+# the last chunk is unequal (12288 vs 3072 for Q/K/V).
+QKV_DIMS = [3072, 3072, 3072]
+LINEAR1_DIMS = [3072, 3072, 3072, 12288]
+
+# Static (non-fused) renames from underscore-flat Flux-layout paths to
+# underscore-flat diffusers paths. Built per-call by format()-ing the layer index.
+DOUBLE_RENAME_TEMPLATES = {
+ 'double_blocks_{i}_img_attn_proj': 'transformer_blocks_{i}_attn_to_out_0',
+ 'double_blocks_{i}_img_mlp_0': 'transformer_blocks_{i}_ff_net_0_proj',
+ 'double_blocks_{i}_img_mlp_2': 'transformer_blocks_{i}_ff_net_2',
+ 'double_blocks_{i}_txt_attn_proj': 'transformer_blocks_{i}_attn_to_add_out',
+ 'double_blocks_{i}_txt_mlp_0': 'transformer_blocks_{i}_ff_context_net_0_proj',
+ 'double_blocks_{i}_txt_mlp_2': 'transformer_blocks_{i}_ff_context_net_2',
+}
+SINGLE_RENAME_TEMPLATES = {
+ 'single_blocks_{i}_linear2': 'single_transformer_blocks_{i}_proj_out',
+}
+
+# Fused-target qkv mappings. Double-block qkv fans out to img-side (to_*) and
+# context-side (add_*_proj). Single-block linear1 fans out to single attn and
+# proj_mlp.
+DOUBLE_IMG_QKV_TARGETS = ('attn_to_q', 'attn_to_k', 'attn_to_v')
+DOUBLE_TXT_QKV_TARGETS = ('attn_add_q_proj', 'attn_add_k_proj', 'attn_add_v_proj')
+SINGLE_LINEAR1_TARGETS = ('attn_to_q', 'attn_to_k', 'attn_to_v', 'proj_mlp')
+
+
+def build_static_rename(num_double, num_single):
+ """Return {flux_flat_name: diffusers_flat_name} for non-fused paths."""
+ out = {}
+ for i in range(num_double):
+ for src, dst in DOUBLE_RENAME_TEMPLATES.items():
+ out[src.format(i=i)] = dst.format(i=i)
+ for i in range(num_single):
+ for src, dst in SINGLE_RENAME_TEMPLATES.items():
+ out[src.format(i=i)] = dst.format(i=i)
+ return out
+
+
+def get_block_counts():
+ """Read num_layers / num_single_layers from the live transformer, with fallback."""
+ sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+ transformer = getattr(sd_model, 'transformer', None)
+ config = getattr(transformer, 'config', None)
+ num_double = getattr(config, 'num_layers', DEFAULT_NUM_DOUBLE_LAYERS) if config is not None else DEFAULT_NUM_DOUBLE_LAYERS
+ num_single = getattr(config, 'num_single_layers', DEFAULT_NUM_SINGLE_LAYERS) if config is not None else DEFAULT_NUM_SINGLE_LAYERS
+ return num_double, num_single
+
+
+class NetworkModuleLokrSliceChunk(network_lokr.NetworkModuleLokr):
+ """LoKR module that returns one row-range of the Kronecker product.
+
+ Used when a LoKR adapter targets a fused weight with unequal chunk sizes
+ (e.g., Chroma single ``linear1`` = Q/K/V/proj_mlp at dims [3072, 3072, 3072, 12288]).
+ The shared ``NetworkModuleLokrChunk`` only supports equal-sized chunks via
+ ``torch.chunk``; this variant slices an explicit row range.
+ """
+ def __init__(self, net, weights, start_row, end_row):
+ super().__init__(net, weights)
+ self.start_row = start_row
+ self.end_row = end_row
+
+ def calc_updown(self, target):
+ if self.w1 is not None:
+ w1 = self.w1.to(target.device, dtype=target.dtype)
+ else:
+ w1a = self.w1a.to(target.device, dtype=target.dtype)
+ w1b = self.w1b.to(target.device, dtype=target.dtype)
+ w1 = w1a @ w1b
+ if self.w2 is not None:
+ w2 = self.w2.to(target.device, dtype=target.dtype)
+ else:
+ from modules.lora import lyco_helpers
+ if self.t2 is None:
+ w2a = self.w2a.to(target.device, dtype=target.dtype)
+ w2b = self.w2b.to(target.device, dtype=target.dtype)
+ w2 = w2a @ w2b
+ else:
+ t2 = self.t2.to(target.device, dtype=target.dtype)
+ w2a = self.w2a.to(target.device, dtype=target.dtype)
+ w2b = self.w2b.to(target.device, dtype=target.dtype)
+ w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+ full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
+ updown = network_lokr.make_kron(full_shape, w1, w2)
+ updown = updown[self.start_row:self.end_row]
+ output_shape = list(updown.shape)
+ return self.finalize_updown(updown, target, output_shape)
+
+
+def try_load_lora(name, network_on_disk, lora_scale):
+ """Try loading a Chroma LoRA (plus DoRA) as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LORA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ static_rename = build_static_rename(*get_block_counts())
+
+ groups = group_by_suffixes(state_dict, LORA_SUFFIXES)
+ groups = expand_chroma_fused_lora(groups)
+ groups = apply_static_rename(groups, static_rename)
+
+ unmapped = 0
+ shape_mismatch = 0
+ for network_key, w in groups.items():
+ if 'lora_down.weight' not in w or 'lora_up.weight' not in w:
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ if not shapes_match(sd_module, w['lora_down.weight'], w['lora_up.weight']):
+ log.warning(f'Network load: type=LoRA name="{name}" key={network_key} shape mismatch')
+ shape_mismatch += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
+
+ return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=shape_mismatch)
+
+
+def try_load_lokr(name, network_on_disk, lora_scale):
+ """Try loading a Chroma LoKR as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOKR_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ static_rename = build_static_rename(*get_block_counts())
+
+ groups = group_by_suffixes(state_dict, LOKR_SUFFIXES)
+ groups, slice_info = expand_chroma_fused_lokr(groups)
+ groups = apply_static_rename(groups, static_rename)
+ slice_info = {static_rename.get(k, k): v for k, v in slice_info.items()}
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w)
+ has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w)
+ if not (has_1 and has_2):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ rng = slice_info.get(network_key)
+ if rng is not None:
+ start, end = rng
+ net.modules[network_key] = NetworkModuleLokrSliceChunk(net, nw, start, end)
+ else:
+ net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw)
+
+ return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped)
+
+
+def try_load_loha(name, network_on_disk, lora_scale):
+ """Try loading a Chroma LoHA as native modules. Fused qkv/linear1 groups are skipped."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOHA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ static_rename = build_static_rename(*get_block_counts())
+
+ groups = group_by_suffixes(state_dict, LOHA_SUFFIXES)
+ groups, skipped = drop_chroma_fused_groups(groups, family='LoHA', name=name)
+ groups = apply_static_rename(groups, static_rename)
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ if not all(k in w for k in ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b")):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_hada.NetworkModuleHada(net, nw)
+
+ return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def try_load_oft(name, network_on_disk, lora_scale):
+ """Try loading a Chroma OFT adapter as native modules. Fused qkv/linear1 groups are skipped."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, OFT_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ static_rename = build_static_rename(*get_block_counts())
+
+ groups = group_by_suffixes(state_dict, OFT_SUFFIXES)
+ groups, skipped = drop_chroma_fused_groups(groups, family='OFT', name=name)
+ groups = apply_static_rename(groups, static_rename)
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ if not ("oft_blocks" in w or "oft_diag" in w):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw)
+
+ return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def has_marker(state_dict, markers):
+ return any(any(m in k for m in markers) for k in state_dict)
+
+
+def resolve_mapping():
+ sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+ lora_convert.assign_network_names_to_compvis_modules(sd_model)
+ return getattr(shared.sd_model, 'network_layer_mapping', {}) or {}
+
+
+def new_network(name, network_on_disk):
+ net = network.Network(name, network_on_disk)
+ net.mtime = os.path.getmtime(network_on_disk.filename)
+ return net
+
+
+def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0, skipped=0):
+ if len(net.modules) == 0:
+ if unmapped or mismatch or skipped:
+ log.debug(
+ f'Network load: type={family} name="{name}" native no-match'
+ f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}'
+ )
+ return None
+ log.debug(
+ f'Network load: type={family} name="{name}" native modules={len(net.modules)}'
+ f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} scale={lora_scale}'
+ )
+ l.timer.activate += time.time() - t0
+ return net
+
+
+def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
+ if not hasattr(sd_module, 'weight'):
+ return False
+ if hasattr(sd_module, 'sdnq_dequantizer'):
+ mod_shape = sd_module.sdnq_dequantizer.original_shape
+ else:
+ mod_shape = sd_module.weight.shape
+ if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2:
+ return False
+ return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0]
+
+
+def group_by_suffixes(state_dict, suffixes):
+ """Group state_dict entries by target module.
+
+ Returns ``{flat_key: {suffix: tensor, ...}}`` where the flat key follows
+ the Flux pre-rename layout (``double_blocks_{i}_img_attn_qkv`` etc.).
+ """
+ groups: dict[str, dict[str, torch.Tensor]] = {}
+ for key, value in state_dict.items():
+ parsed = parse_key(key, suffixes)
+ if parsed is None:
+ continue
+ flat_key, suffix = parsed
+ slot = groups.get(flat_key)
+ if slot is None:
+ slot = {}
+ groups[flat_key] = slot
+ slot[suffix] = value
+ return groups
+
+
+def parse_key(key, suffixes):
+ """Strip prefix and suffix, return (flat_key, normalized_suffix) or None."""
+ stripped = key
+ for p in KNOWN_PREFIXES:
+ if key.startswith(p):
+ stripped = key[len(p):]
+ break
+
+ matched_suffix = None
+ split_at = -1
+ for marker in suffixes:
+ if stripped.endswith(marker):
+ split_at = len(stripped) - len(marker)
+ matched_suffix = marker.lstrip('.')
+ break
+ if split_at < 0:
+ return None
+
+ base = stripped[:split_at]
+ if not base:
+ return None
+
+ suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix)
+ flat_key = base.replace('.', '_')
+ return flat_key, suffix
+
+
+def expand_chroma_fused_lora(groups):
+ """Split fused QKV / linear1 LoRA groups into their per-projection components.
+
+ Chroma LoRAs are trained against Flux's fused-attention layout, while the
+ diffusers ``ChromaTransformer2DModel`` exposes split modules. The
+ fused-attention LoRA convention shares ``down`` across the splits and
+ concatenates ``up`` along dim 0; the inverse here splits ``up`` by the
+ per-target dim list while copying ``down`` to each child.
+ """
+ out: dict[str, dict[str, torch.Tensor]] = {}
+ for key, w in groups.items():
+ if key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv'):
+ stem = key[:-len('img_attn_qkv')] if key.endswith('_img_attn_qkv') else key[:-len('txt_attn_qkv')]
+ targets = DOUBLE_IMG_QKV_TARGETS if key.endswith('_img_attn_qkv') else DOUBLE_TXT_QKV_TARGETS
+ block_prefix = stem.replace('double_blocks_', 'transformer_blocks_')
+ split_groups = split_fused_lora_group(w, QKV_DIMS, [block_prefix + t for t in targets])
+ if split_groups is None:
+ out[key] = w
+ continue
+ out.update(split_groups)
+ elif key.endswith('_linear1') and 'single_blocks' in key:
+ stem = key[:-len('linear1')]
+ block_prefix = stem.replace('single_blocks_', 'single_transformer_blocks_')
+ split_groups = split_fused_lora_group(w, LINEAR1_DIMS, [block_prefix + t for t in SINGLE_LINEAR1_TARGETS])
+ if split_groups is None:
+ out[key] = w
+ continue
+ out.update(split_groups)
+ else:
+ out[key] = w
+ return out
+
+
+def split_fused_lora_group(w, dims, target_keys):
+ """Split a fused LoRA (down, up) into per-target groups by row dim.
+
+ Returns ``{target_key: {suffix: tensor, ...}}`` or ``None`` if the input is
+ malformed (missing tensors, up-weight rows don't sum to dims).
+ """
+ down = w.get('lora_down.weight')
+ up = w.get('lora_up.weight')
+ if down is None or up is None:
+ return None
+ if up.shape[0] != sum(dims):
+ return None
+ alpha = w.get('alpha')
+ dora = w.get('dora_scale')
+ bias = w.get('bias')
+ scale = w.get('scale')
+ out: dict[str, dict[str, torch.Tensor]] = {}
+ start = 0
+ for tk, d in zip(target_keys, dims):
+ chunk_up = up[start:start + d].contiguous()
+ start += d
+ slot = {'lora_down.weight': down, 'lora_up.weight': chunk_up}
+ if alpha is not None:
+ slot['alpha'] = alpha
+ if dora is not None:
+ slot['dora_scale'] = dora
+ if bias is not None:
+ slot['bias'] = bias
+ if scale is not None:
+ slot['scale'] = scale
+ out[tk] = slot
+ return out
+
+
+def expand_chroma_fused_lokr(groups):
+ """Mark fused QKV / linear1 LoKR groups as slice-chunked.
+
+ LoKR factorizations don't compose with row-splitting at load time without
+ materializing the full Kronecker product. Instead, each target gets a
+ shallow copy of the same tensor dict, plus an entry in ``slice_info`` that
+ drives :class:`NetworkModuleLokrSliceChunk` to slice rows lazily on each
+ forward pass.
+ """
+ out: dict[str, dict[str, torch.Tensor]] = {}
+ slice_info: dict[str, tuple[int, int]] = {}
+ for key, w in groups.items():
+ if key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv'):
+ stem = key[:-len('img_attn_qkv')] if key.endswith('_img_attn_qkv') else key[:-len('txt_attn_qkv')]
+ targets = DOUBLE_IMG_QKV_TARGETS if key.endswith('_img_attn_qkv') else DOUBLE_TXT_QKV_TARGETS
+ block_prefix = stem.replace('double_blocks_', 'transformer_blocks_')
+ assign_lokr_slices(out, slice_info, w, QKV_DIMS, [block_prefix + t for t in targets])
+ elif key.endswith('_linear1') and 'single_blocks' in key:
+ stem = key[:-len('linear1')]
+ block_prefix = stem.replace('single_blocks_', 'single_transformer_blocks_')
+ assign_lokr_slices(out, slice_info, w, LINEAR1_DIMS, [block_prefix + t for t in SINGLE_LINEAR1_TARGETS])
+ else:
+ out[key] = w
+ return out, slice_info
+
+
+def assign_lokr_slices(out, slice_info, w, dims, target_keys):
+ start = 0
+ for tk, d in zip(target_keys, dims):
+ out[tk] = dict(w)
+ slice_info[tk] = (start, start + d)
+ start += d
+
+
+def drop_chroma_fused_groups(groups, family, name):
+ """Remove fused QKV / linear1 groups (no chunk variant for LoHA/OFT)."""
+ out: dict[str, dict[str, torch.Tensor]] = {}
+ skipped = 0
+ for key, w in groups.items():
+ is_fused_qkv = key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv')
+ is_fused_linear1 = key.endswith('_linear1') and 'single_blocks' in key
+ if is_fused_qkv or is_fused_linear1:
+ log.warning(f'Network load: type={family} name="{name}" key={key} fused group skipped (unsupported)')
+ skipped += 1
+ continue
+ out[key] = w
+ return out, skipped
+
+
+def apply_static_rename(groups, static_rename):
+ """Rewrite Flux-layout flat keys to diffusers flat keys, then prepend ``lora_transformer_``.
+
+ Keys without an entry in ``static_rename`` are passed through unchanged
+ (they may already be diffusers paths from PEFT-style files, or they may
+ target the ``distilled_guidance_layer`` approximator). The final
+ ``lora_transformer_`` prefix is added uniformly to match the format
+ ``assign_network_names_to_compvis_modules`` registers.
+ """
+ out: dict[str, dict[str, torch.Tensor]] = {}
+ for key, w in groups.items():
+ renamed = static_rename.get(key, key)
+ out['lora_transformer_' + renamed] = w
+ return out
diff --git a/models/Reference/stepfun-ai--Step1X-Edit-v1p2.jpg b/pipelines/ernie/__init__.py
similarity index 100%
rename from models/Reference/stepfun-ai--Step1X-Edit-v1p2.jpg
rename to pipelines/ernie/__init__.py
diff --git a/pipelines/ernie/ernie_image.py b/pipelines/ernie/ernie_image.py
new file mode 100644
index 000000000..f8143e44b
--- /dev/null
+++ b/pipelines/ernie/ernie_image.py
@@ -0,0 +1,196 @@
+"""ERNIE-Image img2img and inpainting pipelines.
+
+Diffusers only ships ``ErnieImagePipeline`` (text-to-image). These two
+subclasses add img2img and inpainting on top of it without touching the
+upstream pipeline.
+
+ERNIE-specific subtleties relative to the Anima equivalent:
+
+- Latent rank is 4-D ``[B, C, H, W]`` (image, not video). No ``unsqueeze(2)``
+ on mask or init latent.
+- The transformer takes a *patched* 128-channel latent at H/16 resolution.
+ The mirror of the upstream decode (BN-unnormalize -> unpatchify -> vae.decode)
+ is ``vae.encode -> patchify -> BN-normalize`` here.
+- Latent normalization uses ``vae.bn.running_mean``/``running_var``
+ applied to the patched 128-channel latent, not scalar
+ ``vae.config.latents_mean``/``latents_std``.
+- ``ErnieImagePipeline`` has no ``image_processor`` instance attribute, so we
+ build a ``VaeImageProcessor`` ad-hoc and cache it on the pipeline.
+"""
+
+from typing import Callable, Dict, List, Optional, Union
+
+import diffusers
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.utils.torch_utils import randn_tensor
+
+from modules import devices
+
+
+def _get_image_processor(pipe):
+ """Build (and cache) a VaeImageProcessor for the pipe.
+
+ pipe.vae_scale_factor=16 already accounts for the 2x2 patchify; the
+ underlying VAE downsamples by 8. The processor needs the VAE-only factor.
+ """
+ proc = getattr(pipe, 'ernie_image_processor', None)
+ if proc is None:
+ proc = VaeImageProcessor(vae_scale_factor=pipe.vae_scale_factor // 2)
+ pipe.ernie_image_processor = proc
+ return proc
+
+
+def _encode_image(pipe, image, dtype, device, generator):
+ """VAE-encode an image and produce the patched, BN-normalized 128-channel latent."""
+ if isinstance(image, list):
+ image = image[0]
+ proc = _get_image_processor(pipe)
+ image_tensor = proc.preprocess(image).to(device=device, dtype=pipe.vae.dtype)
+ raw_latents = pipe.vae.encode(image_tensor).latent_dist.sample(generator)
+ patched = pipe._patchify_latents(raw_latents) # pylint: disable=protected-access
+ bn_mean = pipe.vae.bn.running_mean.to(device=device).view(1, -1, 1, 1)
+ bn_std = torch.sqrt(pipe.vae.bn.running_var.to(device=device).view(1, -1, 1, 1) + 1e-5)
+ return ((patched.float() - bn_mean) / bn_std).to(dtype)
+
+
+def _setup_img2img_schedule(scheduler, strength, num_inference_steps, device):
+ """Set custom sigma schedule, return first sigma after scheduler shift."""
+ custom_sigmas = torch.linspace(max(strength, 0.01), 0.0, num_inference_steps).tolist()
+ scheduler.set_timesteps(sigmas=custom_sigmas, device=device)
+ return scheduler.sigmas[0].item()
+
+
+def _prepare_mask(pipe, mask_image, height, width, device):
+ if isinstance(mask_image, Image.Image):
+ mask_image = mask_image.convert("L")
+ if isinstance(mask_image, Image.Image):
+ import torchvision.transforms.functional as TF
+ mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).to(device=device, dtype=torch.float32)
+ elif isinstance(mask_image, torch.Tensor):
+ mask_tensor = mask_image.to(device=device, dtype=torch.float32)
+ if mask_tensor.ndim == 2:
+ mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
+ elif mask_tensor.ndim == 3:
+ mask_tensor = mask_tensor.unsqueeze(0)
+ else:
+ mask_tensor = torch.ones(1, 1, height, width, device=device, dtype=torch.float32)
+ latent_h = height // pipe.vae_scale_factor
+ latent_w = width // pipe.vae_scale_factor
+ mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest")
+ return mask_latent[:, :1, :, :]
+
+
+class ErnieImageImg2ImgPipeline(diffusers.ErnieImagePipeline):
+ """ERNIE-Image image-to-image pipeline."""
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = "",
+ image: Optional[PipelineImageInput] = None,
+ strength: float = 0.8,
+ height: int = 1024,
+ width: int = 1024,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 4.0,
+ num_images_per_prompt: int = 1,
+ generator: Optional[torch.Generator] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ output_type: str = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ use_pe: bool = True,
+ ):
+ actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+ init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+ noise = randn_tensor(init_latents.shape, generator=generator, device=devices.device, dtype=devices.dtype)
+ noised = actual_sigma * noise + (1.0 - actual_sigma) * init_latents
+
+ orig_set_timesteps = self.scheduler.set_timesteps
+ self.scheduler.set_timesteps = lambda *args, **kwargs: None
+ try:
+ return super().__call__(
+ prompt=prompt, negative_prompt=negative_prompt, height=height, width=width,
+ num_inference_steps=num_inference_steps, guidance_scale=guidance_scale,
+ num_images_per_prompt=num_images_per_prompt, generator=generator, latents=noised,
+ prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds,
+ output_type=output_type, return_dict=return_dict,
+ callback_on_step_end=callback_on_step_end,
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ use_pe=use_pe,
+ )
+ finally:
+ self.scheduler.set_timesteps = orig_set_timesteps
+
+
+class ErnieImageInpaintPipeline(ErnieImageImg2ImgPipeline):
+ """ERNIE-Image inpainting pipeline."""
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = "",
+ image: Optional[PipelineImageInput] = None,
+ mask_image: Optional[PipelineImageInput] = None,
+ strength: float = 0.8,
+ height: int = 1024,
+ width: int = 1024,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 4.0,
+ num_images_per_prompt: int = 1,
+ generator: Optional[torch.Generator] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ output_type: str = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ use_pe: bool = True,
+ ):
+ actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+ init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+ noise = randn_tensor(init_latents.shape, generator=generator, device=devices.device, dtype=devices.dtype)
+ noised = actual_sigma * noise + (1.0 - actual_sigma) * init_latents
+ mask_latent = _prepare_mask(self, mask_image, height, width, devices.device)
+
+ orig_set_timesteps = self.scheduler.set_timesteps
+ self.scheduler.set_timesteps = lambda *args, **kwargs: None
+
+ user_callback = callback_on_step_end
+
+ def blend_callback(pipe, i, t, callback_kwargs):
+ cur_latents = callback_kwargs.get("latents")
+ if cur_latents is not None:
+ sigma_next = pipe.scheduler.sigmas[i + 1].item() if i + 1 < len(pipe.scheduler.sigmas) else 0.0
+ init_at_t = sigma_next * noise + (1.0 - sigma_next) * init_latents
+ mask_dt = mask_latent.to(cur_latents.dtype)
+ blended = mask_dt * cur_latents + (1.0 - mask_dt) * init_at_t.to(cur_latents.dtype)
+ callback_kwargs["latents"] = blended.to(cur_latents.dtype)
+ if user_callback is not None:
+ callback_kwargs = user_callback(pipe, i, t, callback_kwargs)
+ return callback_kwargs
+
+ try:
+ return diffusers.ErnieImagePipeline.__call__(
+ self,
+ prompt=prompt, negative_prompt=negative_prompt,
+ height=height, width=width, num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale, num_images_per_prompt=num_images_per_prompt,
+ generator=generator, latents=noised, prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds, output_type=output_type,
+ return_dict=return_dict, callback_on_step_end=blend_callback,
+ callback_on_step_end_tensor_inputs=["latents"],
+ use_pe=use_pe,
+ )
+ finally:
+ self.scheduler.set_timesteps = orig_set_timesteps
diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py
new file mode 100644
index 000000000..3401a4bec
--- /dev/null
+++ b/pipelines/ernie/ernie_lora.py
@@ -0,0 +1,273 @@
+"""ERNIE-Image native adapter loader.
+
+Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
+(``lora_force_diffusers`` off and ``ernieimage`` in ``allow_native``). Reads the
+safetensors directly and writes into sdnext's existing
+``network_layer_mapping``, returning a ``Network`` populated with
+``NetworkModule*`` entries that ``network_activate`` will apply. If the
+setting is on, the diffusers PEFT path handles the file instead.
+
+Entry points, one per family:
+
+- LoRA (+ DoRA) via :func:`try_load_lora`
+- LoKR via :func:`try_load_lokr`
+- LoHA via :func:`try_load_loha`
+- OFT via :func:`try_load_oft`
+
+Recognized key prefixes for every family: ``diffusion_model.``,
+``transformer.``, ``lora_unet_``, or bare. Diffusers-PEFT ``lora_A``/``lora_B``
+are normalized to ``lora_down``/``lora_up``.
+
+The ERNIE-Image transformer has separate ``self_attention.to_q``/``to_k``/
+``to_v``/``to_out.0`` linear modules (no fused QKV layout), so no
+chunk/split machinery is needed and all four families are supported uniformly.
+"""
+
+import os
+import time
+import torch
+from modules import shared, sd_models
+from modules.logger import log
+from modules.lora import network, network_lora, network_lokr, network_hada, network_oft, lora_convert
+from modules.lora import lora_common as l
+
+
+KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_")
+
+# Every family also picks up the universal optional keys
+# (alpha, scale, bias, dora_scale) via base NetworkModule.__init__.
+LORA_SUFFIXES = (
+ ".lora_down.weight", ".lora_up.weight",
+ ".lora_A.weight", ".lora_B.weight",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOKR_SUFFIXES = (
+ ".lokr_w1", ".lokr_w2",
+ ".lokr_w1_a", ".lokr_w1_b",
+ ".lokr_w2_a", ".lokr_w2_b",
+ ".lokr_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOHA_SUFFIXES = (
+ ".hada_w1_a", ".hada_w1_b",
+ ".hada_w2_a", ".hada_w2_b",
+ ".hada_t1", ".hada_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+OFT_SUFFIXES = (
+ ".oft_blocks", ".oft_diag",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+
+LORA_MARKERS = (".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight")
+LOKR_MARKERS = (".lokr_w1", ".lokr_w2")
+LOHA_MARKERS = (".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b")
+OFT_MARKERS = (".oft_blocks", ".oft_diag")
+
+SUFFIX_NORMALIZE = {
+ "lora_A.weight": "lora_down.weight",
+ "lora_B.weight": "lora_up.weight",
+}
+
+
+def try_load_lora(name, network_on_disk, lora_scale):
+ """Try loading an ERNIE-Image LoRA (plus DoRA) as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LORA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+
+ groups = group_by_suffixes(state_dict, LORA_SUFFIXES)
+
+ unmapped = 0
+ shape_mismatch = 0
+ for network_key, w in groups.items():
+ if 'lora_down.weight' not in w or 'lora_up.weight' not in w:
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ if not shapes_match(sd_module, w['lora_down.weight'], w['lora_up.weight']):
+ log.warning(f'Network load: type=LoRA name="{name}" key={network_key} shape mismatch')
+ shape_mismatch += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
+
+ return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=shape_mismatch)
+
+
+def try_load_lokr(name, network_on_disk, lora_scale):
+ """Try loading an ERNIE-Image LoKR as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOKR_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+
+ groups = group_by_suffixes(state_dict, LOKR_SUFFIXES)
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w)
+ has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w)
+ if not (has_1 and has_2):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw)
+
+ return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped)
+
+
+def try_load_loha(name, network_on_disk, lora_scale):
+ """Try loading an ERNIE-Image LoHA as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOHA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+
+ groups = group_by_suffixes(state_dict, LOHA_SUFFIXES)
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ if not all(k in w for k in ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b")):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_hada.NetworkModuleHada(net, nw)
+
+ return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped)
+
+
+def try_load_oft(name, network_on_disk, lora_scale):
+ """Try loading an ERNIE-Image OFT adapter as native modules."""
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, OFT_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+
+ groups = group_by_suffixes(state_dict, OFT_SUFFIXES)
+
+ unmapped = 0
+ for network_key, w in groups.items():
+ if not ("oft_blocks" in w or "oft_diag" in w):
+ continue
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw)
+
+ return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped)
+
+
+def has_marker(state_dict, markers):
+ return any(any(m in k for m in markers) for k in state_dict)
+
+
+def resolve_mapping():
+ sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+ lora_convert.assign_network_names_to_compvis_modules(sd_model)
+ return getattr(shared.sd_model, 'network_layer_mapping', {}) or {}
+
+
+def new_network(name, network_on_disk):
+ net = network.Network(name, network_on_disk)
+ net.mtime = os.path.getmtime(network_on_disk.filename)
+ return net
+
+
+def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0):
+ if len(net.modules) == 0:
+ if unmapped or mismatch:
+ log.debug(
+ f'Network load: type={family} name="{name}" native no-match'
+ f' unmapped={unmapped} mismatch={mismatch}'
+ )
+ return None
+ log.debug(
+ f'Network load: type={family} name="{name}" native modules={len(net.modules)}'
+ f' unmapped={unmapped} mismatch={mismatch} scale={lora_scale}'
+ )
+ l.timer.activate += time.time() - t0
+ return net
+
+
+def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
+ if not hasattr(sd_module, 'weight'):
+ return False
+ if hasattr(sd_module, 'sdnq_dequantizer'):
+ mod_shape = sd_module.sdnq_dequantizer.original_shape
+ else:
+ mod_shape = sd_module.weight.shape
+ if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2:
+ return False
+ return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0]
+
+
+def group_by_suffixes(state_dict, suffixes):
+ """Group state_dict entries by target module.
+
+ Returns ``{network_key: {suffix: tensor, ...}}`` where ``network_key`` follows
+ the sdnext convention ``lora_transformer_``. Only keys
+ whose suffix appears in ``suffixes`` are kept; ``lora_A``/``lora_B`` are
+ normalized to ``lora_down``/``lora_up``.
+ """
+ groups: dict[str, dict[str, torch.Tensor]] = {}
+ for key, value in state_dict.items():
+ parsed = parse_key(key, suffixes)
+ if parsed is None:
+ continue
+ network_key, suffix = parsed
+ slot = groups.get(network_key)
+ if slot is None:
+ slot = {}
+ groups[network_key] = slot
+ slot[suffix] = value
+ return groups
+
+
+def parse_key(key, suffixes):
+ stripped = key
+ for p in KNOWN_PREFIXES:
+ if key.startswith(p):
+ stripped = key[len(p):]
+ break
+
+ matched_suffix = None
+ split_at = -1
+ for marker in suffixes:
+ if stripped.endswith(marker):
+ split_at = len(stripped) - len(marker)
+ matched_suffix = marker.lstrip('.')
+ break
+ if split_at < 0:
+ return None
+
+ base = stripped[:split_at]
+ if not base:
+ return None
+
+ suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix)
+ network_key = 'lora_transformer_' + base.replace('.', '_')
+ return network_key, suffix
diff --git a/pipelines/f_lite/f_lite.model.py b/pipelines/f_lite/f_lite.model.py
index e56e6545b..3adbb9b85 100644
--- a/pipelines/f_lite/f_lite.model.py
+++ b/pipelines/f_lite/f_lite.model.py
@@ -233,7 +233,7 @@ class PatchEmbed(nn.Module):
self.patch_size = patch_size
def forward(self, x):
- B, C, H, W = x.shape
+ _B, _C, _H, _W = x.shape
x = self.patch_proj(x)
x = rearrange(x, "b c h w -> b (h w) c")
return x
@@ -380,7 +380,7 @@ class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): #
@apply_forward_hook
def forward(self, x, context, timesteps):
- b, c, h, w = x.shape
+ b, _c, h, w = x.shape
x = self.patch_embed(x) # b, T, d
x = torch.cat([self.register_tokens.repeat(b, 1, 1), x], 1) # b, T + N, d
diff --git a/pipelines/f_lite/model.py b/pipelines/f_lite/model.py
index e56e6545b..3adbb9b85 100644
--- a/pipelines/f_lite/model.py
+++ b/pipelines/f_lite/model.py
@@ -233,7 +233,7 @@ class PatchEmbed(nn.Module):
self.patch_size = patch_size
def forward(self, x):
- B, C, H, W = x.shape
+ _B, _C, _H, _W = x.shape
x = self.patch_proj(x)
x = rearrange(x, "b c h w -> b (h w) c")
return x
@@ -380,7 +380,7 @@ class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): #
@apply_forward_hook
def forward(self, x, context, timesteps):
- b, c, h, w = x.shape
+ b, _c, h, w = x.shape
x = self.patch_embed(x) # b, T, d
x = torch.cat([self.register_tokens.repeat(b, 1, 1), x], 1) # b, T + N, d
diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py
index 81e97ea0c..fcea16474 100644
--- a/pipelines/flux/flux2_lora.py
+++ b/pipelines/flux/flux2_lora.py
@@ -1,11 +1,58 @@
-"""Flux2/Klein-specific LoRA loading.
+"""Flux2/Klein native adapter loader.
-Handles:
-- Kohya-format LoRA via native module loading (lora_unet_ prefix keys)
-- LoKR adapters via native module loading (bypasses diffusers PEFT system)
-- Bare BFL-format keys in state dicts (adds diffusion_model. prefix for converter)
+Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
+(``f2`` in ``allow_native``). Reads the safetensors directly and writes into
+sdnext's existing ``network_layer_mapping``, returning a ``Network`` populated
+with ``NetworkModule*`` entries that ``network_activate`` will apply.
-Installed via apply_patch() during pipeline loading.
+Recognized key prefixes for every family: ``diffusion_model.``,
+``transformer.``, ``lora_unet_``, ``lycoris_``, ``base_model.model.``
+(PEFT save wrapper), bare BFL paths (e.g. ``double_blocks.``), and bare
+diffusers paths (``transformer_blocks.`` / ``single_transformer_blocks.``,
+produced by ``Flux2Transformer2DModel.save_lora_adapter()``). Diffusers-PEFT
+``lora_A``/``lora_B`` are normalized to ``lora_down``/``lora_up`` and a
+``..`` infix between the suffix and ``.weight`` (e.g.
+``.lora_A.default.weight``) is stripped to match the standard suffix table.
+
+BFL/kohya keys are mapped to diffusers paths via ``F2_SINGLE_MAP`` /
+``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks is split into
+three Q/K/V targets at lookup time. PEFT keys are diffusers paths already and
+are returned verbatim with no chunking.
+
+Per-family fused-QKV handling:
+
+- LoRA: load-time chunk of ``lora_up`` along dim 0 (the down-side is shared).
+- LoKR: apply-time slice via :class:`NetworkModuleLokrChunk`, which builds
+ ``kron(w1, w2)`` once and returns the designated row range.
+- LoHA: apply-time slice via :class:`NetworkModuleHadaChunk`, which slices
+ ``w1a``/``w2a`` and computes the partial Hadamard product. Tucker
+ (CP-decomposed) LoHAs are not chunked and are skipped on fused targets.
+- OFT, IA3, GLoRA, Full: no chunk class exists and the math is not row-sliceable
+ without re-deriving per-projection structure. Fused groups are skipped with a
+ warning.
+- Norm: targets 1-D LayerNorm/RMSNorm parameters; never fused.
+
+LyCORIS algorithm coverage relative to upstream
+``KohakuBlueleaf/LyCORIS/lycoris/modules/``:
+
+- Native: LoRA, LoKR, LoHA, OFT, BOFT, IA3, GLoRA, Norm, Full.
+- Saved as standard LoRA: LoCon and DyLoRA. Both ``custom_state_dict``
+ outputs collapse to ``lora_up.weight``/``lora_down.weight``/``alpha``
+ (LoCon bakes its ``scalar`` into ``lora_up``; DyLoRA concats its
+ per-block slabs into a max-rank matrix), so ``try_load_lora`` loads
+ them losslessly relative to upstream's own export.
+- Deferred: TLoRA. The file saves only ``q_layer.weight`` /
+ ``p_layer.weight`` / ``lambda_layer`` / ``alpha``; the base SVD
+ reference (``base_q`` / ``base_p`` / ``base_lambda``) that the
+ delta math subtracts is unsaved by upstream design and the
+ ``sig_type`` selection mode is unrecoverable from the file, so
+ any loader has a silent-correctness gap for ``sig_type != 'principal'``.
+ Files fail cleanly with "not loaded".
+
+Diffusers-PEFT fallback (used when ``lora_force_diffusers`` is on) is preserved
+via :func:`apply_patch`, which monkey-patches ``Flux2LoraLoaderMixin.lora_state_dict``
+to inject the ``diffusion_model.`` prefix for bare-BFL keys and bake kohya
+``.alpha`` scaling into ``lora_down`` weights.
"""
import os
@@ -13,19 +60,114 @@ import time
import torch
from modules import shared, sd_models
from modules.logger import log
-from modules.lora import network, network_lokr, network_lora, lora_convert
+from modules.lora import (
+ network, network_lora, network_lokr, network_hada, network_oft, network_boft,
+ network_ia3, network_glora, network_norm, network_full, lora_convert,
+)
from modules.lora import lora_common as l
-BARE_FLUX_PREFIXES = ("single_blocks.", "double_blocks.", "img_in.", "txt_in.",
- "final_layer.", "time_in.", "single_stream_modulation.",
- "double_stream_modulation_")
+# === Format detection ===
-# BFL -> diffusers module path mapping for Flux2/Klein
+# Prefixes we recognize as the "true" format-identifying prefix on a state-dict
+# key. The PEFT save wrapper ``base_model.model.`` is handled separately as a
+# pre-strip step (see :func:`_unwrap_peft_wrapper`) because it can wrap any of
+# the prefixes below — peft.save_pretrained prepends it indiscriminately.
+#
+# - ``diffusion_model.`` — AI-toolkit / BFL native (e.g. ostris/ai-toolkit)
+# - ``transformer.`` — diffusers PEFT in-memory (e.g. HF DreamBooth scripts)
+# - ``lora_unet_`` — kohya-ss/sd-scripts standard
+# - ``lycoris_`` — LyCORIS-standalone save (e.g. SimpleTuner LoKR);
+# the path under this prefix is an underscore-rendered
+# diffusers path, not a BFL path
+KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_", "lycoris_")
+
+BARE_FLUX_PREFIXES = (
+ "single_blocks.", "double_blocks.", "img_in.", "txt_in.",
+ "final_layer.", "time_in.", "single_stream_modulation.",
+ "double_stream_modulation_",
+)
+
+# Bare diffusers paths (no wrapping prefix) — produced by
+# ``Flux2Transformer2DModel.save_lora_adapter()`` after attaching a PEFT adapter.
+# These are already-diffusers paths and pass through ``resolve_targets`` verbatim.
+BARE_DIFFUSERS_PREFIXES = ("single_transformer_blocks.", "transformer_blocks.")
+BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" # sentinel value for ``parse_key`` return
+
+SUFFIX_NORMALIZE = {
+ "lora_A.weight": "lora_down.weight",
+ "lora_B.weight": "lora_up.weight",
+}
+
+
+# === Family suffix tables (alpha / scale / bias / dora_scale flow into weights.w via base NetworkModule.__init__) ===
+
+LORA_SUFFIXES = (
+ ".lora_down.weight", ".lora_up.weight", ".lora_mid.weight",
+ ".lora_A.weight", ".lora_B.weight",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOKR_SUFFIXES = (
+ ".lokr_w1", ".lokr_w2",
+ ".lokr_w1_a", ".lokr_w1_b",
+ ".lokr_w2_a", ".lokr_w2_b",
+ ".lokr_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+LOHA_SUFFIXES = (
+ ".hada_w1_a", ".hada_w1_b",
+ ".hada_w2_a", ".hada_w2_b",
+ ".hada_t1", ".hada_t2",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+OFT_SUFFIXES = (
+ ".oft_blocks", ".oft_diag",
+ ".alpha", ".dora_scale", ".bias", ".scale",
+)
+IA3_SUFFIXES = (
+ ".weight", ".on_input",
+ ".alpha", ".scale",
+)
+GLORA_SUFFIXES = (
+ ".a1.weight", ".a2.weight",
+ ".b1.weight", ".b2.weight",
+ ".alpha", ".dora_scale", ".scale",
+)
+NORM_SUFFIXES = (
+ ".w_norm", ".b_norm",
+ ".alpha", ".scale",
+)
+FULL_SUFFIXES = (
+ ".diff", ".diff_b",
+ ".alpha", ".scale",
+)
+
+LORA_MARKERS = (
+ ".lora_down.weight", ".lora_up.weight",
+ ".lora_A.weight", ".lora_B.weight",
+ # PEFT named-adapter saves embed the slot name as ``.lora_A..weight``;
+ # the trailing-dot forms catch every variant.
+ ".lora_A.", ".lora_B.",
+)
+LOKR_MARKERS = (".lokr_w1", ".lokr_w2")
+LOHA_MARKERS = (".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b")
+OFT_MARKERS = (".oft_blocks", ".oft_diag")
+IA3_MARKERS = (".on_input",) # NOT .weight — too generic, overlaps every other family
+GLORA_MARKERS = (".a1.weight", ".a2.weight", ".b1.weight", ".b2.weight")
+NORM_MARKERS = (".w_norm",)
+FULL_MARKERS = (".diff",)
+
+
+# === BFL → diffusers mapping ===
+
+# Single-block (single_transformer_blocks.{i}.) — both projections are single fused diffusers modules,
+# so no chunking is needed for any adapter family.
F2_SINGLE_MAP = {
'linear1': 'attn.to_qkv_mlp_proj',
'linear2': 'attn.to_out',
}
+
+# Double-block non-QKV targets (transformer_blocks.{i}.).
F2_DOUBLE_MAP = {
'img_attn.proj': 'attn.to_out.0',
'txt_attn.proj': 'attn.to_add_out',
@@ -34,14 +176,16 @@ F2_DOUBLE_MAP = {
'txt_mlp.0': 'ff_context.linear_in',
'txt_mlp.2': 'ff_context.linear_out',
}
+
+# Double-block fused QKV targets — diffusers exposes Q/K/V as separate modules,
+# so resolve_targets emits three (path, chunk_index, num_chunks=3) entries.
F2_QKV_MAP = {
'img_attn.qkv': ('attn', ['to_q', 'to_k', 'to_v']),
'txt_attn.qkv': ('attn', ['add_q_proj', 'add_k_proj', 'add_v_proj']),
}
-
-# Kohya underscore suffix -> BFL dot suffix (last underscore becomes dot)
-# Used to convert kohya key fragments to look up F2_DOUBLE_MAP / F2_QKV_MAP
+# Kohya underscore suffix → BFL dot suffix (last underscore becomes dot).
+# Used to convert kohya key fragments to look up F2_DOUBLE_MAP / F2_QKV_MAP.
KOHYA_SUFFIX_MAP = {
'img_attn_proj': 'img_attn.proj',
'txt_attn_proj': 'txt_attn.proj',
@@ -54,208 +198,194 @@ KOHYA_SUFFIX_MAP = {
}
-def try_load_lora(name, network_on_disk, lora_scale):
- """Try loading a Flux2/Klein LoRA as native modules.
+# === Shared scaffolding ===
- Handles three key formats:
- - Kohya: lora_unet_double_blocks_0_img_attn_proj.lora_down.weight
- - AI toolkit (BFL): diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight
- - Diffusers PEFT: transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_A.weight
- Returns a Network with native modules, or None to fall through to the diffusers path.
- """
- t0 = time.time()
- state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
- has_lora = any('.lora_down.' in k or '.lora_up.' in k or '.lora_A.' in k or '.lora_B.' in k for k in state_dict)
- if not has_lora:
- return None
- is_f2_keys = any(
- k.startswith(('lora_unet_single_blocks_', 'lora_unet_double_blocks_',
- 'diffusion_model.single_blocks.', 'diffusion_model.double_blocks.',
- 'transformer.single_transformer_blocks.', 'transformer.transformer_blocks.'))
- for k in state_dict
- )
- if not is_f2_keys:
- return None
- net = load_lora_native(name, network_on_disk, state_dict)
+def has_marker(state_dict, markers):
+ return any(any(m in k for m in markers) for k in state_dict)
+
+
+def resolve_mapping():
+ sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+ lora_convert.assign_network_names_to_compvis_modules(sd_model)
+ return getattr(shared.sd_model, 'network_layer_mapping', {}) or {}
+
+
+def new_network(name, network_on_disk):
+ net = network.Network(name, network_on_disk)
+ net.mtime = os.path.getmtime(network_on_disk.filename)
+ return net
+
+
+def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0, skipped=0):
if len(net.modules) == 0:
+ if unmapped or mismatch or skipped:
+ log.debug(
+ f'Network load: type={family} name="{name}" native no-match'
+ f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}'
+ )
return None
- log.debug(f'Network load: type=LoRA name="{name}" native modules={len(net.modules)} scale={lora_scale}')
+ log.debug(
+ f'Network load: type={family} name="{name}" native modules={len(net.modules)}'
+ f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} scale={lora_scale}'
+ )
l.timer.activate += time.time() - t0
return net
-def _group_lora_keys(state_dict):
- """Group LoRA state dict keys into (targets, weights_dict) pairs.
+def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
+ if not hasattr(sd_module, 'weight'):
+ return False
+ if hasattr(sd_module, 'sdnq_dequantizer'):
+ mod_shape = sd_module.sdnq_dequantizer.original_shape
+ else:
+ mod_shape = sd_module.weight.shape
+ if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2:
+ return False
+ return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0]
- Normalizes all three formats into a common structure. Weight keys are
- normalized to lora_down.weight / lora_up.weight regardless of input naming.
- Returns list of (targets, weights_dict) where targets come from BFL->diffusers mapping.
+
+def _unwrap_peft_wrapper(key):
+ """Strip the ``base_model.model.`` prefix added by ``peft.save_pretrained``.
+
+ PeftModel.save_pretrained prepends this wrapper to every adapter key. The
+ content underneath can be any of the formats KNOWN_PREFIXES already handle:
+
+ - BFL keys (e.g. fal/flux-2-klein-4B-outpaint-lora:
+ ``base_model.model.double_blocks.0.img_attn.proj.lora_A.weight``)
+ - Diffusers paths under ``transformer.`` (HF DreamBooth scripts that
+ target diffusers modules and let peft wrap them)
+ - Bare-BFL keys (rare but possible)
+
+ Stripping the wrapper once is enough; the rest of :func:`parse_key` then
+ matches the unwrapped key against KNOWN_PREFIXES or the bare-BFL fallback
+ normally. Mirrors the diffusers ``Flux2LoraLoaderMixin.lora_state_dict``
+ behavior at lora_pipeline.py:5684-5686, which renames the prefix to
+ ``diffusion_model.`` before feeding the key to the AI-toolkit converter.
"""
- # Detect format from first relevant key
- sample = next((k for k in state_dict if '.lora_' in k), None)
- if sample is None:
- return []
+ if key.startswith("base_model.model."):
+ return key[len("base_model.model."):]
+ return key
- if sample.startswith('lora_unet_'):
- return _group_kohya(state_dict)
- elif sample.startswith('diffusion_model.'):
- return _group_bfl(state_dict)
- elif sample.startswith('transformer.'):
- return _group_peft(state_dict)
- # Bare BFL keys (no prefix)
- if any(k.startswith(p) for k in state_dict for p in BARE_FLUX_PREFIXES):
- return _group_bfl(state_dict, prefix='')
+def _strip_peft_adapter_name(key):
+ """Normalize ``.lora_[AB]..weight`` to ``.lora_[AB].weight``.
+ ``peft.PeftModel`` and the diffusers ``save_lora_adapter`` exporter embed the
+ adapter slot name into the saved key (``"default"`` when not explicitly
+ set). Strip a single non-dotted name segment so the suffix table matches
+ without listing every plausible adapter name.
+ """
+ for inner in (".lora_A.", ".lora_B."):
+ idx = key.find(inner)
+ if idx == -1:
+ continue
+ rest = key[idx + len(inner):]
+ if rest == "weight" or not rest.endswith(".weight"):
+ continue
+ adapter_name = rest[:-len(".weight")]
+ if adapter_name and "." not in adapter_name:
+ return key[:idx] + inner + "weight"
+ return key
+
+
+def parse_key(key, suffixes):
+ """Return ``(prefix_used, base, suffix_normalized)`` or ``None``.
+
+ ``prefix_used`` is the matched ``KNOWN_PREFIXES`` element, or ``None`` for
+ bare BFL keys. ``base`` is the format-native module path (kohya / lycoris
+ underscore-style or BFL / diffusers dot-style depending on prefix).
+ """
+ key = _unwrap_peft_wrapper(key)
+ key = _strip_peft_adapter_name(key)
+ prefix_used = None
+ stripped = key
+ for p in KNOWN_PREFIXES:
+ if key.startswith(p):
+ prefix_used = p
+ stripped = key[len(p):]
+ break
+ if prefix_used is None:
+ if any(key.startswith(p) for p in BARE_DIFFUSERS_PREFIXES):
+ prefix_used = BARE_DIFFUSERS_PREFIX_USED
+ elif not any(key.startswith(p) for p in BARE_FLUX_PREFIXES):
+ return None
+
+ matched_suffix = None
+ split_at = -1
+ for marker in suffixes:
+ if stripped.endswith(marker):
+ split_at = len(stripped) - len(marker)
+ matched_suffix = marker.lstrip('.')
+ break
+ if split_at < 0:
+ return None
+
+ base = stripped[:split_at]
+ if not base:
+ return None
+
+ suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix)
+ return prefix_used, base, suffix
+
+
+def group_by_suffixes(state_dict, suffixes):
+ """Group state_dict entries by ``(prefix_used, base)``.
+
+ Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where
+ ``prefix_used`` is a ``KNOWN_PREFIXES`` element or ``None`` for bare-BFL.
+ Per-family loaders apply their own key-presence gates on each group.
+ """
+ groups: dict[tuple, dict[str, torch.Tensor]] = {}
+ for key, value in state_dict.items():
+ parsed = parse_key(key, suffixes)
+ if parsed is None:
+ continue
+ prefix_used, base, suffix = parsed
+ slot = groups.get((prefix_used, base))
+ if slot is None:
+ slot = {}
+ groups[(prefix_used, base)] = slot
+ slot[suffix] = value
+ return groups
+
+
+def resolve_targets(prefix_used, base):
+ """Return ``[(diffusers_path, chunk_index, num_chunks), ...]`` for a parsed group key.
+
+ For kohya prefix, applies ``KOHYA_SUFFIX_MAP`` then ``F2_*_MAP``. For
+ BFL/bare-BFL, applies ``F2_*_MAP`` directly. For PEFT (``transformer.``),
+ returns the base verbatim with no chunking — it is already a diffusers path.
+ """
+ if prefix_used == 'lora_unet_':
+ return _kohya_to_diffusers_targets(base)
+ if prefix_used in (None, 'diffusion_model.'):
+ return _bfl_to_diffusers_targets(base)
+ if prefix_used == 'transformer.':
+ return [(base, None, None)]
+ if prefix_used == BARE_DIFFUSERS_PREFIX_USED:
+ # Already-diffusers path with no wrapping prefix (e.g. produced by
+ # Flux2Transformer2DModel.save_lora_adapter()). Pass through verbatim.
+ return [(base, None, None)]
+ if prefix_used == 'lycoris_':
+ # base is an already-underscored diffusers path (e.g.
+ # 'transformer_blocks_0_attn_add_k_proj'). The caller's network_key
+ # construction does base.replace('.', '_'); for already-underscored
+ # paths that's a no-op, so the network_key matches the entry stamped
+ # by lora_convert.assign_network_names_to_compvis_modules
+ # (e.g. 'lora_transformer_transformer_blocks_0_attn_add_k_proj').
+ return [(base, None, None)]
return []
-def _normalize_weight_key(suffix):
- """lora_A.weight -> lora_down.weight, lora_B.weight -> lora_up.weight"""
- return suffix.replace('lora_A.', 'lora_down.').replace('lora_B.', 'lora_up.')
-
-
-def _group_kohya(state_dict):
- """Group kohya-format keys (lora_unet_ prefix, underscored module names)."""
- groups = {}
- for key, weight in state_dict.items():
- if not key.startswith('lora_unet_'):
- continue
- base, _, suffix = key.partition('.')
- if not suffix:
- continue
- if base not in groups:
- groups[base] = {}
- groups[base][_normalize_weight_key(suffix)] = weight
-
- results = []
- for base, weights_dict in groups.items():
- if 'lora_down.weight' not in weights_dict:
- continue
- stripped = base[len('lora_unet_'):]
- targets = _kohya_key_to_targets(stripped)
- if targets:
- results.append((targets, weights_dict))
- return results
-
-
-def _group_bfl(state_dict, prefix='diffusion_model.'):
- """Group BFL/AI-toolkit-format keys (dot-separated module names)."""
- groups = {}
- for key, weight in state_dict.items():
- if prefix and not key.startswith(prefix):
- continue
- stripped = key[len(prefix):]
- # Split at lora boundary: double_blocks.0.img_attn.proj.lora_A.weight
- for marker in ('.lora_A.', '.lora_B.', '.lora_down.', '.lora_up.', '.alpha', '.dora_scale'):
- pos = stripped.find(marker)
- if pos != -1:
- base = stripped[:pos]
- suffix = stripped[pos + 1:] if stripped[pos + 1:] else marker[1:] # handle bare .dora_scale / .alpha
- break
- else:
- continue
- if base not in groups:
- groups[base] = {}
- groups[base][_normalize_weight_key(suffix)] = weight
-
- results = []
- for base, weights_dict in groups.items():
- if 'lora_down.weight' not in weights_dict:
- continue
- targets = _bfl_key_to_targets(base)
- if targets:
- results.append((targets, weights_dict))
- return results
-
-
-def _group_peft(state_dict):
- """Group diffusers PEFT-format keys (transformer. prefix, diffusers module names)."""
- groups = {}
- for key, weight in state_dict.items():
- if not key.startswith('transformer.'):
- continue
- stripped = key[len('transformer.'):]
- for marker in ('.lora_A.', '.lora_B.', '.lora_down.', '.lora_up.', '.alpha'):
- pos = stripped.find(marker)
- if pos != -1:
- module_path = stripped[:pos]
- suffix = stripped[pos + 1:]
- break
- else:
- continue
- if module_path not in groups:
- groups[module_path] = {}
- groups[module_path][_normalize_weight_key(suffix)] = weight
-
- results = []
- for module_path, weights_dict in groups.items():
- if 'lora_down.weight' not in weights_dict:
- continue
- # Already in diffusers path format — direct target, no mapping needed
- results.append(([(module_path, None, None)], weights_dict))
- return results
-
-
-def load_lora_native(name, network_on_disk, state_dict):
- """Load Flux2/Klein LoRA as native modules from any supported key format."""
- sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
- lora_convert.assign_network_names_to_compvis_modules(sd_model)
- net = network.Network(name, network_on_disk)
- net.mtime = os.path.getmtime(network_on_disk.filename)
-
- for targets, weights_dict in _group_lora_keys(state_dict):
- for module_path, chunk_index, num_chunks in targets:
- network_key = "lora_transformer_" + module_path.replace(".", "_")
- sd_module = sd_model.network_layer_mapping.get(network_key)
- if sd_module is None:
- continue
-
- w = {}
- if chunk_index is not None:
- up = weights_dict['lora_up.weight']
- chunks = torch.chunk(up, num_chunks, dim=0)
- w['lora_up.weight'] = chunks[chunk_index].contiguous()
- w['lora_down.weight'] = weights_dict['lora_down.weight']
- else:
- w['lora_up.weight'] = weights_dict['lora_up.weight']
- w['lora_down.weight'] = weights_dict['lora_down.weight']
-
- # Validate dimensions match the target module
- if hasattr(sd_module, 'weight'):
- if hasattr(sd_module, 'sdnq_dequantizer'):
- mod_shape = sd_module.sdnq_dequantizer.original_shape
- else:
- mod_shape = sd_module.weight.shape
- if w['lora_down.weight'].shape[1] != mod_shape[1] or w['lora_up.weight'].shape[0] != mod_shape[0]:
- log.warning(f'Network load: type=LoRA shape mismatch: {network_key} lora={w["lora_down.weight"].shape[1]}x{w["lora_up.weight"].shape[0]} module={mod_shape[1]}x{mod_shape[0]}')
- continue
-
- if 'alpha' in weights_dict:
- w['alpha'] = weights_dict['alpha']
- if 'dora_scale' in weights_dict:
- w['dora_scale'] = weights_dict['dora_scale']
-
- nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
- net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
-
- return net
-
-
-def _kohya_key_to_targets(stripped):
- """Map a stripped kohya key to (diffusers_module_path, chunk_index, num_chunks) targets.
-
- Input examples: 'double_blocks_0_img_attn_proj', 'single_blocks_5_linear1'
- """
- targets = []
-
+def _kohya_to_diffusers_targets(stripped):
+ """For kohya keys like ``double_blocks_0_img_attn_proj`` or ``single_blocks_5_linear1``."""
+ targets: list[tuple[str, int | None, int | None]] = []
if stripped.startswith('single_blocks_'):
rest = stripped[len('single_blocks_'):]
idx, _, suffix = rest.partition('_')
if suffix in F2_SINGLE_MAP:
targets.append((f'single_transformer_blocks.{idx}.{F2_SINGLE_MAP[suffix]}', None, None))
-
elif stripped.startswith('double_blocks_'):
rest = stripped[len('double_blocks_'):]
idx, _, kohya_suffix = rest.partition('_')
@@ -268,46 +398,452 @@ def _kohya_key_to_targets(stripped):
attn_prefix, proj_keys = F2_QKV_MAP[bfl_suffix]
for i, proj_key in enumerate(proj_keys):
targets.append((f'transformer_blocks.{idx}.{attn_prefix}.{proj_key}', i, len(proj_keys)))
-
return targets
-def _bfl_key_to_targets(base):
- """Map a BFL dot-separated key to (diffusers_module_path, chunk_index, num_chunks) targets.
-
- Input examples: 'double_blocks.0.img_attn.proj', 'single_blocks.5.linear1'
- Same mapping as LoKR uses.
- """
- targets = []
+def _bfl_to_diffusers_targets(base):
+ """For BFL keys like ``double_blocks.0.img_attn.proj`` or ``single_blocks.5.linear1``."""
+ targets: list[tuple[str, int | None, int | None]] = []
parts = base.split('.')
if len(parts) < 3:
return targets
-
block_type, block_idx, module_suffix = parts[0], parts[1], '.'.join(parts[2:])
-
if block_type == 'single_blocks' and module_suffix in F2_SINGLE_MAP:
- path = f'single_transformer_blocks.{block_idx}.{F2_SINGLE_MAP[module_suffix]}'
- targets.append((path, None, None))
+ targets.append((f'single_transformer_blocks.{block_idx}.{F2_SINGLE_MAP[module_suffix]}', None, None))
elif block_type == 'double_blocks':
if module_suffix in F2_DOUBLE_MAP:
- path = f'transformer_blocks.{block_idx}.{F2_DOUBLE_MAP[module_suffix]}'
- targets.append((path, None, None))
+ targets.append((f'transformer_blocks.{block_idx}.{F2_DOUBLE_MAP[module_suffix]}', None, None))
elif module_suffix in F2_QKV_MAP:
attn_prefix, proj_keys = F2_QKV_MAP[module_suffix]
for i, proj_key in enumerate(proj_keys):
- path = f'transformer_blocks.{block_idx}.{attn_prefix}.{proj_key}'
- targets.append((path, i, len(proj_keys)))
-
+ targets.append((f'transformer_blocks.{block_idx}.{attn_prefix}.{proj_key}', i, len(proj_keys)))
return targets
-def apply_lora_alphas(state_dict):
- """Bake kohya-format .alpha scaling into lora_down weights and remove alpha keys.
+# === Native loaders ===
- Diffusers' Flux2 converter only handles lora_A/lora_B (or lora_down/lora_up) keys.
- Kohya-format LoRAs store per-layer alpha values as separate .alpha keys that the
- converter doesn't consume, causing a ValueError on leftover keys. This matches the
- approach used by _convert_kohya_flux_lora_to_diffusers for Flux 1.
+
+def try_load(name, network_on_disk, lora_scale):
+ """Run every Flux2 family loader in dispatch order, merge any that match.
+
+ Per-family ``try_load_*`` entry points stay public; this is the single
+ umbrella the dispatcher in ``modules.lora.lora_load.load_safetensors``
+ calls. Order matters only for marker-cost: LoRA / LoKR are most common
+ so their fast bail-out runs first; the rare families come last.
+
+ Returns a ``Network`` with the union of modules from every matching
+ family loader, or ``None`` if no loader recognized the file.
+ """
+ net = None
+ for try_fn in (
+ try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
+ try_load_ia3, try_load_glora, try_load_norm, try_load_full,
+ ):
+ sub = try_fn(name, network_on_disk, lora_scale)
+ if sub is None:
+ continue
+ if net is None:
+ net = sub
+ else:
+ net.modules.update(sub.modules)
+ return net
+
+
+def try_load_lora(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein LoRA (plus DoRA via the universal ``finalize_updown`` hook) as native modules.
+
+ Handles kohya, AI-toolkit/BFL, diffusers PEFT, and bare-BFL key formats.
+ Fused QKV in double_blocks is split at load time by chunking the up-weight
+ along dim 0; the down-weight is shared across Q/K/V.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LORA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, LORA_SUFFIXES)
+
+ unmapped = 0
+ mismatch = 0
+ for (prefix, base), w in groups.items():
+ if 'lora_down.weight' not in w or 'lora_up.weight' not in w:
+ continue
+ for diffusers_path, chunk_idx, num_chunks in resolve_targets(prefix, base):
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+
+ if chunk_idx is not None:
+ chunks = torch.chunk(w['lora_up.weight'], num_chunks, dim=0)
+ target_w = dict(w)
+ target_w['lora_up.weight'] = chunks[chunk_idx].contiguous()
+ else:
+ target_w = w
+
+ if not shapes_match(sd_module, target_w['lora_down.weight'], target_w['lora_up.weight']):
+ log.warning(
+ f'Network load: type=LoRA name="{name}" key={network_key}'
+ f' lora={target_w["lora_down.weight"].shape[1]}x{target_w["lora_up.weight"].shape[0]}'
+ f' module={getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}'
+ f' shape mismatch'
+ )
+ mismatch += 1
+ continue
+
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module)
+ net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
+
+ return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=mismatch)
+
+
+def try_load_lokr(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein LoKR as native modules.
+
+ Stores only the compact LoKR factors (``w1``/``w2``) and computes
+ ``kron(w1, w2)`` on-the-fly during weight application. For fused QKV
+ targets in double_blocks, :class:`NetworkModuleLokrChunk` materializes the
+ full Kronecker product and returns the designated Q/K/V slice.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOKR_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, LOKR_SUFFIXES)
+
+ unmapped = 0
+ for (prefix, base), w in groups.items():
+ has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w)
+ has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w)
+ if not (has_1 and has_2):
+ continue
+ for diffusers_path, chunk_idx, num_chunks in resolve_targets(prefix, base):
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ if chunk_idx is not None:
+ net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, nw, chunk_idx, num_chunks)
+ else:
+ net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw)
+
+ return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped)
+
+
+def try_load_loha(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein LoHA (Hadamard product) adapter as native modules.
+
+ Standard non-Tucker LoHA on fused QKV in double_blocks is supported via
+ :class:`NetworkModuleHadaChunk`, which slices ``w1a``/``w2a`` at the
+ chunk's row range and computes the partial Hadamard. Tucker
+ (CP-decomposed) LoHAs are skipped on fused targets because the chunk
+ class does not implement the CP path; non-fused Tucker LoHAs go through
+ the standard :class:`NetworkModuleHada`.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, LOHA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, LOHA_SUFFIXES)
+
+ unmapped = 0
+ skipped = 0
+ for (prefix, base), w in groups.items():
+ if not all(k in w for k in ('hada_w1_a', 'hada_w1_b', 'hada_w2_a', 'hada_w2_b')):
+ continue
+ is_tucker = 'hada_t1' in w or 'hada_t2' in w
+ targets = resolve_targets(prefix, base)
+ is_fused = any(t[1] is not None for t in targets)
+ if is_fused and is_tucker:
+ log.warning(f'Network load: type=LoHA name="{name}" key={base} Tucker fused QKV skipped (unsupported)')
+ skipped += 1
+ continue
+ for diffusers_path, chunk_idx, num_chunks in targets:
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ if chunk_idx is not None:
+ net.modules[network_key] = network_hada.NetworkModuleHadaChunk(net, nw, chunk_idx, num_chunks)
+ else:
+ net.modules[network_key] = network_hada.NetworkModuleHada(net, nw)
+
+ return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def try_load_oft(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein OFT or BOFT adapter as native modules.
+
+ Both algorithms share the ``oft_blocks`` save key and are discriminated
+ by tensor dimensionality, mirroring LyCORIS's own ``algo_check``:
+
+ - **OFT** — 3-D ``(num_blocks, block_size, block_size)``. Both kohya
+ (``oft_blocks`` + alpha-as-constraint) and LyCORIS (``oft_diag``)
+ layouts route through :class:`NetworkModuleOFT`.
+ - **BOFT** — 4-D ``(boft_m, block_num, block_size, block_size)``,
+ a cascade of butterfly factors. Routes through
+ :class:`NetworkModuleBOFT` which ports the butterfly-cascade
+ ``make_weight`` from LyCORIS boft.py.
+
+ Fused QKV in double_blocks is skipped with a warning for both: an OFT
+ block structure (and BOFT's per-stage block partition) is tied to the
+ target module's ``out_features``, so a per-Q/K/V split would require
+ re-deriving the rotations per chunk. Single-block ``linear1`` (a single
+ fused diffusers module) and all non-QKV double-block targets work fully.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, OFT_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, OFT_SUFFIXES)
+
+ unmapped = 0
+ skipped = 0
+ for (prefix, base), w in groups.items():
+ if not ('oft_blocks' in w or 'oft_diag' in w):
+ continue
+ is_boft = 'oft_blocks' in w and w['oft_blocks'].ndim == 4
+ targets = resolve_targets(prefix, base)
+ if any(t[1] is not None for t in targets):
+ log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" key={base} fused QKV skipped (unsupported)')
+ skipped += 1
+ continue
+ for diffusers_path, _, _ in targets:
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ if is_boft:
+ net.modules[network_key] = network_boft.NetworkModuleBOFT(net, nw)
+ else:
+ net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw)
+
+ return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def try_load_ia3(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein IA3 adapter as native modules.
+
+ IA3 stores a per-row or per-column scale vector keyed under ``.weight``
+ plus an ``.on_input`` flag selecting which axis. The ``.on_input`` marker
+ is the format disambiguator — ``.weight`` alone is too generic and
+ overlaps every other family's ``.lora_down.weight`` / ``.hada_w*`` keys,
+ so the SUFFIXES table includes it but the MARKERS gate insists on
+ ``.on_input``.
+
+ Fused QKV in double_blocks is skipped: ``on_input=True`` IA3 vectors
+ would replicate cleanly to Q/K/V (same ``in_features``) but
+ ``on_input=False`` requires slicing the output-axis vector across the
+ three projections, and there is zero real-world IA3-on-DiT prevalence to
+ justify the asymmetry.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, IA3_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, IA3_SUFFIXES)
+
+ unmapped = 0
+ skipped = 0
+ for (prefix, base), w in groups.items():
+ if not ('weight' in w and 'on_input' in w):
+ continue
+ targets = resolve_targets(prefix, base)
+ if any(t[1] is not None for t in targets):
+ log.warning(f'Network load: type=IA3 name="{name}" key={base} fused QKV skipped (unsupported)')
+ skipped += 1
+ continue
+ for diffusers_path, _, _ in targets:
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_ia3.NetworkModuleIa3(net, nw)
+
+ return finalize_network(net, name, 'IA3', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def try_load_glora(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein GLoRA adapter as native modules.
+
+ GLoRA stores four low-rank components (``a1``/``a2``/``b1``/``b2``) and
+ computes ``ΔW = w2b @ w1b + (target @ w2a) @ w1a`` — the second term is
+ target-dependent. Fused QKV in double_blocks is skipped with a warning
+ because the target-dependent term doesn't slice cleanly without
+ redirecting calc_updown to a fused proxy weight, and zero real-world
+ GLoRA-on-DiT files exist.
+
+ Depends on the ``self.dim`` initialization fix in network_glora.py so
+ that alpha-based ``calc_scale`` is honored.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, GLORA_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, GLORA_SUFFIXES)
+
+ unmapped = 0
+ skipped = 0
+ for (prefix, base), w in groups.items():
+ if not all(k in w for k in ('a1.weight', 'a2.weight', 'b1.weight', 'b2.weight')):
+ continue
+ targets = resolve_targets(prefix, base)
+ if any(t[1] is not None for t in targets):
+ log.warning(f'Network load: type=GLoRA name="{name}" key={base} fused QKV skipped (unsupported)')
+ skipped += 1
+ continue
+ for diffusers_path, _, _ in targets:
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_glora.NetworkModuleGLora(net, nw)
+
+ return finalize_network(net, name, 'GLoRA', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+def try_load_norm(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein Norm adapter (LayerNorm/RMSNorm weight + bias deltas) as native modules.
+
+ Norm adapters target the RMSNorm modules inside Flux2 attention
+ (``attn.norm_q``, ``attn.norm_k``, ``attn.norm_added_q``,
+ ``attn.norm_added_k``) — the only norm modules in Flux2 with trainable
+ weights. The block-level ``norm1``/``norm2`` LayerNorms have
+ ``elementwise_affine=False`` and are not adaptable.
+
+ Loader-local stamping: ``modules/lora/lora_convert.py:assign_network_names_to_compvis_modules``
+ deliberately skips setting ``module.network_layer_name`` for transformer
+ norm modules (except SD3) because of legacy CompVis UNet collisions. This
+ loader bypasses the guard locally — for each target it actually binds, it
+ sets ``network_layer_name`` directly on the host module so
+ ``network_activate`` will apply the delta. No edit to the shared
+ ``lora_convert`` carve-out is required, and no norm module is touched
+ unless a Norm adapter explicitly targets it.
+
+ BFL/kohya prefix support is deferred — there is no public Flux2 BFL norm
+ mapping table to verify against. PEFT prefix (the format produced by
+ ``peft`` training) works directly because the base path is already a
+ diffusers path.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, NORM_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, NORM_SUFFIXES)
+
+ unmapped = 0
+ for (prefix, base), w in groups.items():
+ if 'w_norm' not in w:
+ continue
+ targets = resolve_targets(prefix, base)
+ if not targets:
+ unmapped += 1
+ continue
+ for diffusers_path, chunk_idx, _ in targets:
+ if chunk_idx is not None:
+ continue # norm targets are not fused
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ # Bypass the lora_convert.py:502 transformer-norm guard locally.
+ # Stamping is idempotent and only touches modules a Norm adapter targets.
+ if not getattr(sd_module, 'network_layer_name', None):
+ sd_module.network_layer_name = network_key
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_norm.NetworkModuleNorm(net, nw)
+
+ return finalize_network(net, name, 'Norm', lora_scale, t0, unmapped=unmapped)
+
+
+def try_load_full(name, network_on_disk, lora_scale):
+ """Load a Flux2/Klein Full (full-rank) adapter as native modules.
+
+ Full adapters carry a complete weight delta (``diff``, same shape as the
+ host weight) and an optional bias delta (``diff_b``) via
+ :class:`NetworkModuleFull`. Most realistic use: small per-block bias-only
+ adjustments in distillation LoRAs.
+
+ Fused QKV in double_blocks is skipped with a warning. Full's ``diff`` has
+ the host weight's full shape; row-slicing across three projections is
+ well-defined arithmetically but no chunk class exists and zero
+ real-world Full-on-fused-DiT files exist. Single-block linear1 (a single
+ fused diffusers module) and non-QKV double-block targets work fully.
+ """
+ t0 = time.time()
+ state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
+ if not has_marker(state_dict, FULL_MARKERS):
+ return None
+
+ mapping = resolve_mapping()
+ net = new_network(name, network_on_disk)
+ groups = group_by_suffixes(state_dict, FULL_SUFFIXES)
+
+ unmapped = 0
+ skipped = 0
+ for (prefix, base), w in groups.items():
+ if 'diff' not in w:
+ continue
+ targets = resolve_targets(prefix, base)
+ if any(t[1] is not None for t in targets):
+ log.warning(f'Network load: type=Full name="{name}" key={base} fused QKV skipped (unsupported)')
+ skipped += 1
+ continue
+ for diffusers_path, _, _ in targets:
+ network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
+ sd_module = mapping.get(network_key)
+ if sd_module is None:
+ unmapped += 1
+ continue
+ nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
+ net.modules[network_key] = network_full.NetworkModuleFull(net, nw)
+
+ return finalize_network(net, name, 'Full', lora_scale, t0, unmapped=unmapped, skipped=skipped)
+
+
+# === Diffusers-PEFT path helpers (used when lora_force_diffusers is on) ===
+
+
+def apply_lora_alphas(state_dict):
+ """Bake kohya-format ``.alpha`` scaling into ``lora_down`` weights and remove alpha keys.
+
+ Diffusers' Flux2 converter only handles ``lora_A``/``lora_B`` (or
+ ``lora_down``/``lora_up``) keys. Kohya-format LoRAs store per-layer alpha
+ values as separate ``.alpha`` keys that the converter does not consume,
+ causing a ``ValueError`` on leftover keys. This matches the approach used
+ by ``_convert_kohya_flux_lora_to_diffusers`` for Flux 1.
"""
alpha_keys = [k for k in state_dict if k.endswith('.alpha')]
if not alpha_keys:
@@ -339,8 +875,8 @@ def apply_lora_alphas(state_dict):
def preprocess_f2_keys(state_dict):
- """Add 'diffusion_model.' prefix to bare BFL-format keys so
- Flux2LoraLoaderMixin's format detection routes them to the converter."""
+ """Add ``diffusion_model.`` prefix to bare BFL-format keys so
+ ``Flux2LoraLoaderMixin``'s format detection routes them to the converter."""
if any(k.startswith("diffusion_model.") or k.startswith("base_model.model.") for k in state_dict):
return state_dict
if any(k.startswith(p) for k in state_dict for p in BARE_FLUX_PREFIXES):
@@ -349,98 +885,16 @@ def preprocess_f2_keys(state_dict):
return state_dict
-def try_load_lokr(name, network_on_disk, lora_scale):
- """Try loading a Flux2/Klein LoRA as LoKR native modules.
-
- Returns a Network with native modules if the state dict contains LoKR keys,
- or None to fall through to the generic diffusers path.
- """
- t0 = time.time()
- state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
- if not any('.lokr_w1' in k for k in state_dict):
- return None
- net = load_lokr_native(name, network_on_disk, state_dict)
- if len(net.modules) == 0:
- log.error(f'Network load: type=LoKR name="{name}" no modules matched')
- return None
- log.debug(f'Network load: type=LoKR name="{name}" native modules={len(net.modules)} scale={lora_scale}')
- l.timer.activate += time.time() - t0
- return net
-
-
-def load_lokr_native(name, network_on_disk, state_dict):
- """Load Flux2 LoKR as native modules applied at inference time.
-
- Stores only the compact LoKR factors (w1, w2) and computes kron(w1, w2)
- on-the-fly during weight application. For fused QKV modules in double
- blocks, NetworkModuleLokrChunk computes the full Kronecker product and
- returns only its designated Q/K/V chunk, then frees the temporary.
- """
- prefix = "diffusion_model."
- sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
- lora_convert.assign_network_names_to_compvis_modules(sd_model)
- net = network.Network(name, network_on_disk)
- net.mtime = os.path.getmtime(network_on_disk.filename)
-
- for key in list(state_dict.keys()):
- if not key.endswith('.lokr_w1'):
- continue
- if not key.startswith(prefix):
- continue
-
- base = key[len(prefix):].rsplit('.lokr_w1', 1)[0]
- lokr_weights = {}
- for suffix in ['lokr_w1', 'lokr_w2', 'lokr_w1_a', 'lokr_w1_b', 'lokr_w2_a', 'lokr_w2_b', 'lokr_t2', 'alpha']:
- full_key = f'{prefix}{base}.{suffix}'
- if full_key in state_dict:
- lokr_weights[suffix] = state_dict[full_key]
-
- parts = base.split('.')
- block_type, block_idx, module_suffix = parts[0], parts[1], '.'.join(parts[2:])
-
- targets = [] # (module_path, chunk_index, num_chunks)
- if block_type == 'single_blocks' and module_suffix in F2_SINGLE_MAP:
- path = f'single_transformer_blocks.{block_idx}.{F2_SINGLE_MAP[module_suffix]}'
- targets.append((path, None, None))
- elif block_type == 'double_blocks':
- if module_suffix in F2_DOUBLE_MAP:
- path = f'transformer_blocks.{block_idx}.{F2_DOUBLE_MAP[module_suffix]}'
- targets.append((path, None, None))
- elif module_suffix in F2_QKV_MAP:
- attn_prefix, proj_keys = F2_QKV_MAP[module_suffix]
- for i, proj_key in enumerate(proj_keys):
- path = f'transformer_blocks.{block_idx}.{attn_prefix}.{proj_key}'
- targets.append((path, i, len(proj_keys)))
-
- for module_path, chunk_index, num_chunks in targets:
- network_key = "lora_transformer_" + module_path.replace(".", "_")
- sd_module = sd_model.network_layer_mapping.get(network_key)
- if sd_module is None:
- log.warning(f'Network load: type=LoKR module not found in mapping: {network_key}')
- continue
- weights = network.NetworkWeights(
- network_key=network_key,
- sd_key=network_key,
- w=dict(lokr_weights),
- sd_module=sd_module,
- )
- if chunk_index is not None:
- net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, weights, chunk_index, num_chunks)
- else:
- net.modules[network_key] = network_lokr.NetworkModuleLokr(net, weights)
-
- return net
-
-
patched = False
def apply_patch():
- """Patch Flux2LoraLoaderMixin.lora_state_dict to handle bare BFL-format keys.
+ """Patch ``Flux2LoraLoaderMixin.lora_state_dict`` to handle bare BFL-format keys.
- When a LoRA file has bare BFL keys (no diffusion_model. prefix), the original
- lora_state_dict won't detect them as AI toolkit format. This patch checks for
- bare keys after the original returns and adds the prefix + re-runs conversion.
+ When a LoRA file has bare BFL keys (no ``diffusion_model.`` prefix), the
+ original ``lora_state_dict`` won't detect them as AI toolkit format. This
+ patch checks for bare keys after the original returns and adds the prefix +
+ re-runs conversion. Used only on the diffusers-PEFT fallback path.
"""
global patched # pylint: disable=global-statement
if patched:
diff --git a/pipelines/generic.py b/pipelines/generic.py
index 63196fe08..b65592905 100644
--- a/pipelines/generic.py
+++ b/pipelines/generic.py
@@ -20,6 +20,8 @@ def _loader(component):
def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None):
+ if shared.state.interrupted:
+ return None
transformer = None
if load_config is None:
load_config = {}
@@ -101,6 +103,8 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None):
+ if shared.state.interrupted:
+ return None
text_encoder = None
if load_config is None:
load_config = {}
@@ -250,3 +254,35 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
devices.torch_gc()
shared.state.end(jobid)
return text_encoder
+
+
+def load_vae_override(pipe, load_config=None, override_cls=None, override_args={}):
+ if shared.state.interrupted:
+ return
+ if (shared.opts.sd_vae in [None, 'None', 'Default', 'Automatic']):
+ return
+ if (pipe is None) or (getattr(pipe, 'vae', None) is None):
+ return
+ if load_config is None:
+ load_config = {}
+
+ cls = override_cls or pipe.vae.__class__
+ if not hasattr(cls, 'from_single_file'):
+ log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} safetensors=unsupported')
+ return
+ load_args, quant_args = model_quant.get_dit_args(load_config, module='VAE')
+ log.info(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} args={load_args} quant={quant_args}')
+ try:
+ fn = os.path.join(shared.opts.vae_dir, shared.opts.sd_vae)
+ vae = cls.from_single_file(
+ fn,
+ cache_dir=shared.opts.hfcache_dir,
+ **override_args,
+ **load_args,
+ **quant_args,
+ )
+ if vae is not None:
+ pipe.vae = vae
+ except Exception as e:
+ log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} {e}')
+ # errors.display(e, 'Load')
diff --git a/pipelines/hidream/pipeline_hidream_image_editing.py b/pipelines/hidream/hidream_e1.py
similarity index 100%
rename from pipelines/hidream/pipeline_hidream_image_editing.py
rename to pipelines/hidream/hidream_e1.py
diff --git a/pipelines/hidream/hidream_o1.py b/pipelines/hidream/hidream_o1.py
new file mode 100644
index 000000000..5000f48de
--- /dev/null
+++ b/pipelines/hidream/hidream_o1.py
@@ -0,0 +1,458 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from tqdm.rich import tqdm
+
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+from pipelines.hidream.scheduler_flashfloweuler import FlashFlowMatchEulerDiscreteScheduler
+from pipelines.hidream.scheduler_flowunipc import FlowUniPCMultistepScheduler
+from pipelines.hidream.hidream_o1_utils import (
+ CONDITION_IMAGE_SIZE,
+ DEFAULT_TIMESTEPS,
+ NOISE_SCALE,
+ PATCH_SIZE,
+ TIMESTEP_TOKEN_NUM,
+ T_EPS,
+ _calculate_dimensions,
+ _ensure_special_tokens,
+ _image_to_patch_tensor,
+ _patches_to_np,
+ _pil_to_normalized_tensor,
+ _resize_pilimage,
+ build_t2i_text_sample,
+ get_rope_index_fix_point,
+)
+
+use_flash_attn = False
+try:
+ import flash_attn # pylint: disable=unused-import
+ use_flash_attn = True
+except ImportError:
+ pass
+
+
+@dataclass
+class HiDreamO1PipelineOutput(BaseOutput):
+ images: Union[List[PIL.Image.Image], np.ndarray]
+
+
+class HiDreamO1Pipeline(DiffusionPipeline):
+ model_cpu_offload_seq = "transformer"
+ _callback_tensor_inputs = ["latents"]
+ vae_scale_factor = 1
+
+ def __init__(
+ self,
+ transformer,
+ processor,
+ tokenizer,
+ scheduler,
+ ):
+ super().__init__()
+ self.register_modules(transformer=transformer, processor=processor, scheduler=scheduler, tokenizer=tokenizer)
+ _ensure_special_tokens(self.tokenizer)
+
+ def _build_scheduler(
+ self,
+ num_inference_steps: int,
+ shift: float,
+ device: torch.device,
+ ):
+ if num_inference_steps <= 28:
+ self.scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
+ timesteps_list = DEFAULT_TIMESTEPS if num_inference_steps == 28 else None
+ else:
+ self.scheduler = FlowUniPCMultistepScheduler(use_dynamic_shifting=False, shift=shift)
+ timesteps_list = None
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
+ if timesteps_list is not None:
+ self.scheduler.timesteps = torch.tensor(timesteps_list, device=device, dtype=torch.long)
+ sigmas = [t.item() / 1000.0 for t in self.scheduler.timesteps]
+ sigmas.append(0.0)
+ self.scheduler.sigmas = torch.tensor(sigmas, device=device)
+
+ def _prepare_reference_paths(
+ self,
+ image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]],
+ ref_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]],
+ ) -> List[PIL.Image.Image]:
+ refs: List[PIL.Image.Image] = []
+ if image is not None:
+ if isinstance(image, list):
+ refs.extend(image)
+ else:
+ refs.append(image)
+ if ref_images is not None:
+ if isinstance(ref_images, list):
+ refs.extend(ref_images)
+ else:
+ refs.append(ref_images)
+ return refs
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: str,
+ negative_prompt: Optional[str] = None,
+ image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None,
+ ref_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None,
+ height: int = 1440,
+ width: int = 2560,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 5.0,
+ shift: float = 3.0,
+ timesteps_list: Optional[List[int]] = None,
+ scheduler: Optional[Union[FlashFlowMatchEulerDiscreteScheduler, FlowUniPCMultistepScheduler]] = None,
+ generator: Optional[torch.Generator] = None,
+ seed: Optional[int] = None,
+ noise_scale_start: float = NOISE_SCALE,
+ noise_scale_end: float = NOISE_SCALE,
+ noise_clip_std: float = 0.0,
+ keep_original_aspect: bool = True,
+ callback_on_step_end: Optional[Callable[[DiffusionPipeline, int, int, Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]] = None,
+ callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
+ callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback_steps: int = 1,
+ output_type: str = "pil",
+ return_dict: bool = True,
+ **kwargs,
+ ) -> Union[HiDreamO1PipelineOutput, tuple]:
+ model = self.transformer
+ processor = self.processor
+ tokenizer = self.tokenizer
+ model_config = model.config
+
+ if isinstance(prompt, list):
+ prompt = prompt[0] if len(prompt) > 0 else ""
+ if isinstance(negative_prompt, list):
+ negative_prompt = negative_prompt[0] if len(negative_prompt) > 0 else ""
+ if num_inference_steps <= 28:
+ guidance_scale = 1.0
+
+ device = self._execution_device
+ try:
+ dtype = next(model.parameters()).dtype
+ except (StopIteration, AttributeError, TypeError):
+ dtype = torch.bfloat16
+
+ if callback_on_step_end_tensor_inputs is not None:
+ invalid_inputs = [name for name in callback_on_step_end_tensor_inputs if name not in self._callback_tensor_inputs]
+ if invalid_inputs:
+ raise ValueError(
+ f"callback_on_step_end_tensor_inputs has to be in {self._callback_tensor_inputs}, but found {invalid_inputs}"
+ )
+
+ refs = [img.convert("RGB") for img in self._prepare_reference_paths(image, ref_images)]
+ preresized_ref_pil = None
+
+ if keep_original_aspect and len(refs) >= 1:
+ preresized_ref_pil = _resize_pilimage(refs[0], 2048, PATCH_SIZE)
+ width, height = preresized_ref_pil.size
+ else:
+ width = max(PATCH_SIZE, int(round(width / PATCH_SIZE)) * PATCH_SIZE)
+ height = max(PATCH_SIZE, int(round(height / PATCH_SIZE)) * PATCH_SIZE)
+
+ h_patches = height // PATCH_SIZE
+ w_patches = width // PATCH_SIZE
+
+ if len(refs) == 0:
+ cond_sample = build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config)
+ uncond_sample = None
+ if guidance_scale > 1.0:
+ uncond_sample = build_t2i_text_sample(negative_prompt or " ", height, width, tokenizer, processor, model_config)
+
+ def to_device(sample):
+ return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in sample.items()}
+
+ cond_sample = to_device(cond_sample)
+ if uncond_sample is not None:
+ uncond_sample = to_device(uncond_sample)
+
+ ref_patches = None
+ tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
+ samples = [cond_sample]
+ if uncond_sample is not None:
+ samples.append(uncond_sample)
+ else:
+ image_token_id = model_config.image_token_id
+ video_token_id = model_config.video_token_id
+ vision_start_token_id = model_config.vision_start_token_id
+ spatial_merge_size = model_config.vision_config.spatial_merge_size
+
+ ref_pils = [preresized_ref_pil] if preresized_ref_pil is not None else refs
+ k_refs = len(ref_pils)
+
+ max_size = max(height, width)
+ if k_refs == 2:
+ max_size = max_size * 48 // 64
+ elif k_refs <= 4:
+ max_size = max_size // 2
+ elif k_refs <= 8:
+ max_size = max_size * 24 // 64
+ elif k_refs > 8:
+ max_size = max_size // 4
+
+ ref_pils_resized, ref_patches_list = [], []
+ for pil in ref_pils:
+ pil_r = pil if (preresized_ref_pil is not None and pil is preresized_ref_pil) else _resize_pilimage(pil, max_size, PATCH_SIZE)
+ ref_pils_resized.append(pil_r)
+ x = _pil_to_normalized_tensor(pil_r).unsqueeze(0)
+ x = _image_to_patch_tensor(x, patch_size=PATCH_SIZE).squeeze(0)
+ ref_patches_list.append(x)
+
+ ref_image_lens = [img.shape[0] for img in ref_patches_list]
+ total_ref_len = sum(ref_image_lens)
+ ref_patches = torch.cat(ref_patches_list, dim=0).unsqueeze(0).to(device, dtype)
+
+ tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
+
+ cond_img_size = CONDITION_IMAGE_SIZE
+ if k_refs > 4 and k_refs <= 8:
+ cond_img_size = CONDITION_IMAGE_SIZE * 48 // 64
+ elif k_refs > 8:
+ cond_img_size = CONDITION_IMAGE_SIZE // 2
+
+ ref_pils_vlm = []
+ for pil_r in ref_pils_resized:
+ cond_w, cond_h = _calculate_dimensions(cond_img_size, pil_r.width / pil_r.height)
+ ref_pils_vlm.append(pil_r.resize((cond_w, cond_h), resample=PIL.Image.Resampling.LANCZOS))
+
+ image_grid_thw_tgt = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0)
+ image_grid_thw_ref = torch.zeros((k_refs, 3), dtype=torch.int64)
+ for i, pil_r in enumerate(ref_pils_resized):
+ rw, rh = pil_r.size
+ image_grid_thw_ref[i] = torch.tensor([1, rh // PATCH_SIZE, rw // PATCH_SIZE], dtype=torch.int64)
+
+ samples = []
+ captions = [prompt]
+ if guidance_scale > 1.0:
+ captions.append(negative_prompt or " ")
+
+ for caption in captions:
+ boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
+ tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
+
+ content = [{"type": "image"} for _ in range(k_refs)]
+ content.append({"type": "text", "text": caption})
+ messages = [{"role": "user", "content": content}]
+
+ template_caption = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ proc = processor(text=[template_caption], images=ref_pils_vlm, padding="longest", return_tensors="pt")
+
+ input_ids_2 = tokenizer.encode(boi_token + tms_token * TIMESTEP_TOKEN_NUM, return_tensors="pt", add_special_tokens=False)
+ input_ids = torch.cat([proc.input_ids, input_ids_2], dim=-1)
+
+ igthw_cond = proc.image_grid_thw.clone()
+ for i in range(k_refs):
+ igthw_cond[i, 1] //= spatial_merge_size
+ igthw_cond[i, 2] //= spatial_merge_size
+ igthw_all = torch.cat([igthw_cond, image_grid_thw_tgt, image_grid_thw_ref], dim=0)
+
+ vision_tokens_list = []
+ vt_tgt = torch.full((1, tgt_image_len), image_token_id, dtype=input_ids.dtype)
+ vt_tgt[0, 0] = vision_start_token_id
+ vision_tokens_list.append(vt_tgt)
+ for ref_len in ref_image_lens:
+ vt_ref = torch.full((1, ref_len), image_token_id, dtype=input_ids.dtype)
+ vt_ref[0, 0] = vision_start_token_id
+ vision_tokens_list.append(vt_ref)
+ vision_tokens = torch.cat(vision_tokens_list, dim=1)
+ input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
+
+ position_ids, _ = get_rope_index_fix_point(
+ 1,
+ image_token_id,
+ video_token_id,
+ vision_start_token_id,
+ input_ids=input_ids_pad,
+ image_grid_thw=igthw_all,
+ video_grid_thw=None,
+ attention_mask=None,
+ skip_vision_start_token=[0] * k_refs + [1] + [1] * k_refs,
+ )
+
+ txt_seq_len = input_ids.shape[-1]
+ all_seq_len = position_ids.shape[-1]
+
+ token_types_raw = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
+ bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
+ end = bgn + tgt_image_len + TIMESTEP_TOKEN_NUM
+ token_types_raw[0, bgn:end] = 1
+ token_types_raw[0, end : end + total_ref_len] = 2
+ token_types_raw[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3
+
+ vinput_mask = torch.logical_or(token_types_raw == 1, token_types_raw == 2)
+ token_types_bin = (token_types_raw > 0).to(token_types_raw.dtype)
+
+ samples.append(
+ {
+ "input_ids": input_ids_pad.to(device),
+ "position_ids": position_ids.to(device),
+ "token_types": token_types_bin.to(device),
+ "vinput_mask": vinput_mask.to(device),
+ "pixel_values": proc.pixel_values.to(device, dtype),
+ "image_grid_thw": proc.image_grid_thw.to(device),
+ }
+ )
+
+ if generator is None:
+ generator = torch.Generator()
+ if seed is not None:
+ generator.manual_seed(seed + 1)
+
+ noise = noise_scale_start * randn_tensor((1, 3, height, width), generator=generator, device=device, dtype=dtype)
+ z = _image_to_patch_tensor(noise, patch_size=PATCH_SIZE)
+
+
+ if scheduler is None:
+ self._build_scheduler(
+ num_inference_steps=num_inference_steps,
+ shift=shift,
+ device=device,
+ )
+
+ num_steps = len(self.scheduler.timesteps)
+ if num_steps > 1:
+ noise_scale_schedule = [noise_scale_start + (noise_scale_end - noise_scale_start) * i / (num_steps - 1) for i in range(num_steps)]
+ else:
+ noise_scale_schedule = [noise_scale_start]
+
+ def forward_once(sample: Dict[str, torch.Tensor], z_in: torch.Tensor, t_pixeldit: torch.Tensor):
+ kwargs: Dict[str, Any] = {
+ "input_ids": sample["input_ids"],
+ "position_ids": sample["position_ids"],
+ "vinputs": z_in,
+ "timestep": t_pixeldit.reshape(-1).to(device),
+ "token_types": sample["token_types"],
+ }
+ if use_flash_attn is not None:
+ kwargs["use_flash_attn"] = use_flash_attn
+ if "pixel_values" in sample:
+ kwargs["pixel_values"] = sample["pixel_values"]
+ if "image_grid_thw" in sample:
+ kwargs["image_grid_thw"] = sample["image_grid_thw"]
+
+ outputs = model(**kwargs)
+ x_pred = outputs.x_pred
+ if ref_patches is None:
+ return x_pred[0, sample["vinput_mask"][0]].unsqueeze(0)
+ return x_pred[0, sample["vinput_mask"][0]][:tgt_image_len].unsqueeze(0)
+
+ preview_x0 = None
+ for step_idx, step_t in enumerate(tqdm(self.scheduler.timesteps, desc="Processing", unit="it")):
+ t_pixeldit = 1.0 - step_t.float() / 1000.0
+ sigma = (step_t.float() / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS)
+
+ if ref_patches is None:
+ x_pred_cond = forward_once(samples[0], z.clone(), t_pixeldit)
+ v_cond = (x_pred_cond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma
+ if len(samples) > 1:
+ x_pred_uncond = forward_once(samples[1], z.clone(), t_pixeldit)
+ v_uncond = (x_pred_uncond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma
+ v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
+ else:
+ v_guided = v_cond
+ preview_x0 = x_pred_cond
+ else:
+ vinputs = torch.cat([z, ref_patches], dim=1)
+ x_vis_list = [forward_once(sample, vinputs, t_pixeldit) for sample in samples]
+ x_vis_stacked = torch.cat(x_vis_list, dim=0)
+
+ z_rep = z.expand(len(samples), -1, -1)
+ v_pred = (x_vis_stacked.to(dtype=torch.float32) - z_rep.to(dtype=torch.float32)) / sigma
+ v_cond = v_pred[0:1]
+ if len(samples) > 1:
+ v_uncond = v_pred[1:]
+ v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
+ else:
+ v_guided = v_cond
+ preview_x0 = x_vis_list[0]
+
+ model_output = -v_guided
+ if num_inference_steps <= 28:
+ z = self.scheduler.step(
+ model_output.float(),
+ step_t.to(dtype=torch.float32),
+ z.float(),
+ s_noise=noise_scale_schedule[step_idx],
+ noise_clip_std=noise_clip_std,
+ return_dict=False,
+ )[0].to(dtype)
+ else:
+ z = self.scheduler.step(model_output.float(), step_t.to(dtype=torch.float32), z.float(), return_dict=False)[0].to(dtype)
+
+ if callback_on_step_end is not None:
+ pil_image = PIL.Image.fromarray(_patches_to_np(preview_x0, h_patches, w_patches, invert=False))
+ callback_on_step_end(self, step_idx, int(step_t.item()), {"image": pil_image})
+
+ if output_type == "pil":
+ output_images = [PIL.Image.fromarray(_patches_to_np(preview_x0, h_patches, w_patches, invert=False, rescale=False))]
+ elif output_type == "np":
+ output_images = [_patches_to_np(preview_x0, h_patches, w_patches, invert=True, rescale=True)]
+ else:
+ raise ValueError(f"Unsupported output_type={output_type!r}; supported values are 'pil' and 'np'")
+
+ if not return_dict:
+ return (output_images,)
+ return HiDreamO1PipelineOutput(images=output_images)
+
+
+class HiDreamO1ImagePipeline(HiDreamO1Pipeline):
+ def __call__(
+ self,
+ prompt: str,
+ negative_prompt: Optional[str] = None,
+ image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None,
+ height: int = 1440,
+ width: int = 2560,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 5.0,
+ shift: float = 3.0,
+ timesteps_list: Optional[List[int]] = None,
+ scheduler: Optional[Union[FlashFlowMatchEulerDiscreteScheduler, FlowUniPCMultistepScheduler]] = None,
+ generator: Optional[torch.Generator] = None,
+ seed: Optional[int] = None,
+ noise_scale_start: float = NOISE_SCALE,
+ noise_scale_end: float = NOISE_SCALE,
+ noise_clip_std: float = 0.0,
+ keep_original_aspect: bool = True,
+ use_flash_attn: Optional[bool] = None,
+ callback: Optional[Callable[[int, int, Callable[[], PIL.Image.Image]], None]] = None,
+ output_type: str = "pil",
+ return_dict: bool = True,
+ **kwargs,
+ ) -> Union[HiDreamO1PipelineOutput, tuple]:
+ # image is list, first entry should go to image and remaining to ref_images
+ ref_images = []
+ if isinstance(image, list):
+ ref_images = image[1:]
+ image = image[0] if len(image) > 0 else None
+ return super().__call__(
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ image=image,
+ ref_images=ref_images,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ shift=shift,
+ timesteps_list=timesteps_list,
+ scheduler=scheduler,
+ generator=generator,
+ seed=seed,
+ noise_scale_start=noise_scale_start,
+ noise_scale_end=noise_scale_end,
+ noise_clip_std=noise_clip_std,
+ keep_original_aspect=keep_original_aspect,
+ callback=callback,
+ output_type=output_type,
+ return_dict=return_dict,
+ **kwargs,
+ )
diff --git a/pipelines/hidream/hidream_o1_utils.py b/pipelines/hidream/hidream_o1_utils.py
new file mode 100644
index 000000000..b02e5231a
--- /dev/null
+++ b/pipelines/hidream/hidream_o1_utils.py
@@ -0,0 +1,369 @@
+import math
+from typing import Optional
+
+import numpy as np
+import PIL.Image
+import torch
+
+
+TIMESTEP_TOKEN_NUM = 1
+NOISE_SCALE = 8.0
+T_EPS = 0.001
+CONDITION_IMAGE_SIZE = 384
+PATCH_SIZE = 32
+
+DEFAULT_TIMESTEPS = [
+ 999,
+ 987,
+ 974,
+ 960,
+ 945,
+ 929,
+ 913,
+ 895,
+ 877,
+ 857,
+ 836,
+ 814,
+ 790,
+ 764,
+ 737,
+ 707,
+ 675,
+ 640,
+ 602,
+ 560,
+ 515,
+ 464,
+ 409,
+ 347,
+ 278,
+ 199,
+ 110,
+ 8,
+]
+
+PREDEFINED_RESOLUTIONS = [
+ (2048, 2048),
+ (2304, 1728),
+ (1728, 2304),
+ (2560, 1440),
+ (1440, 2560),
+ (2496, 1664),
+ (1664, 2496),
+ (3104, 1312),
+ (1312, 3104),
+ (2304, 1792),
+ (1792, 2304),
+]
+
+
+def _ensure_special_tokens(tokenizer):
+ if not hasattr(tokenizer, "boi_token"):
+ tokenizer.boi_token = "<|boi_token|>"
+ if not hasattr(tokenizer, "bor_token"):
+ tokenizer.bor_token = "<|bor_token|>"
+ if not hasattr(tokenizer, "eor_token"):
+ tokenizer.eor_token = "<|eor_token|>"
+ if not hasattr(tokenizer, "bot_token"):
+ tokenizer.bot_token = "<|bot_token|>"
+ if not hasattr(tokenizer, "tms_token"):
+ tokenizer.tms_token = "<|tms_token|>"
+
+
+def _find_closest_resolution(width: int, height: int):
+ img_ratio = width / height
+ best_res = PREDEFINED_RESOLUTIONS[0]
+ min_diff = float("inf")
+ for w, h in PREDEFINED_RESOLUTIONS:
+ diff = abs((w / h) - img_ratio)
+ if diff < min_diff:
+ min_diff = diff
+ best_res = (w, h)
+ return best_res
+
+
+def _resize_pilimage(
+ pil_image: PIL.Image.Image,
+ image_size: int,
+ patch_size: int = 16,
+ resampler: PIL.Image.Resampling = PIL.Image.Resampling.BICUBIC,
+):
+ while min(*pil_image.size) >= 2 * image_size:
+ pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=PIL.Image.Resampling.BOX)
+
+ width, height = pil_image.width, pil_image.height
+ max_area = image_size * image_size
+ scale = math.sqrt(max_area / (width * height))
+
+ m = patch_size
+ new_sizes = [
+ (round(width * scale) // m * m, round(height * scale) // m * m),
+ (round(width * scale) // m * m, math.floor(height * scale) // m * m),
+ (math.floor(width * scale) // m * m, round(height * scale) // m * m),
+ (math.floor(width * scale) // m * m, math.floor(height * scale) // m * m),
+ ]
+ new_sizes = sorted(new_sizes, key=lambda x: x[0] * x[1], reverse=True)
+
+ new_size = new_sizes[-1]
+ for candidate in new_sizes:
+ if candidate[0] * candidate[1] <= max_area:
+ new_size = candidate
+ break
+
+ s1 = width / new_size[0]
+ s2 = height / new_size[1]
+ if s1 < s2:
+ pil_image = pil_image.resize([new_size[0], round(height / s1)], resample=resampler)
+ top = (round(height / s1) - new_size[1]) // 2
+ pil_image = pil_image.crop((0, top, new_size[0], top + new_size[1]))
+ else:
+ pil_image = pil_image.resize([round(width / s2), new_size[1]], resample=resampler)
+ left = (round(width / s2) - new_size[0]) // 2
+ pil_image = pil_image.crop((left, 0, left + new_size[0], new_size[1]))
+
+ return pil_image
+
+
+def _calculate_dimensions(max_size: int, ratio: float):
+ width = math.sqrt(max_size * max_size * ratio)
+ height = width / ratio
+ width = int(width / 32) * 32
+ height = int(height / 32) * 32
+ return width, height
+
+
+def _image_to_patch_tensor(x: torch.Tensor, patch_size: int = PATCH_SIZE) -> torch.Tensor:
+ b, c, h, w = x.shape
+ h_patch = h // patch_size
+ w_patch = w // patch_size
+ x = x.reshape(b, c, h_patch, patch_size, w_patch, patch_size)
+ x = x.permute(0, 2, 4, 1, 3, 5)
+ return x.reshape(b, h_patch * w_patch, c * patch_size * patch_size)
+
+
+def _patch_tensor_to_image(x: torch.Tensor, h_patches: int, w_patches: int, patch_size: int = PATCH_SIZE) -> torch.Tensor:
+ b = x.shape[0]
+ c = x.shape[-1] // (patch_size * patch_size)
+ x = x.reshape(b, h_patches, w_patches, c, patch_size, patch_size)
+ x = x.permute(0, 3, 1, 4, 2, 5)
+ return x.reshape(b, c, h_patches * patch_size, w_patches * patch_size)
+
+
+def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor:
+ arr = np.array(image.convert("RGB"), dtype=np.float32) / 255.0
+ x = torch.from_numpy(arr).permute(2, 0, 1)
+ return x * 2.0 - 1.0
+
+
+def _patches_to_np(
+ z: torch.Tensor,
+ h_patches: int,
+ w_patches: int,
+ invert: bool = True,
+ rescale: bool = False,
+) -> np.ndarray:
+ z = z.float()
+ if rescale:
+ clip_ratio = (torch.abs(z) > 1.0).float().mean().item() # hidream-o1 often has outlines and clipping
+ if clip_ratio > 0.10: # Balanced quantiles and headroom: compress outliers while preserving contrast.
+ lo_q, hi_q, headroom = 0.03, 0.97, 1.33
+ elif clip_ratio > 0.03:
+ lo_q, hi_q, headroom = 0.02, 0.98, 1.22
+ else:
+ lo_q, hi_q, headroom = 0.01, 0.99, 1.15
+ z_flat = z.reshape(-1)
+ q_lo = torch.quantile(z_flat, lo_q)
+ q_hi = torch.quantile(z_flat, hi_q)
+ center = (q_hi + q_lo) * 0.5
+ half_range = (q_hi - q_lo) * 0.5
+ if half_range > 0:
+ z = (z - center) / torch.clamp(half_range * headroom, min=0.35) # Balanced robust remap: moderate compression to handle outliers without losing contrast.
+ z = z.clamp(-1.0, 1.0)
+ image = (1.0 - z) / 2.0 if invert else (z + 1.0) / 2.0
+ image = _patch_tensor_to_image(image, h_patches=h_patches, w_patches=w_patches, patch_size=PATCH_SIZE)
+ np_image = image[0].cpu().numpy().transpose(1, 2, 0)
+ np_image = np.round(255.0 * np_image).astype(np.uint8)
+ return np_image
+
+
+def get_rope_index_fix_point(
+ spatial_merge_size,
+ image_token_id,
+ video_token_id,
+ vision_start_token_id,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ skip_vision_start_token=None,
+ fix_point=4096,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if video_grid_thw is not None:
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+ video_grid_thw[:, 0] = 1
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids_i in enumerate(total_input_ids):
+ input_ids_i = input_ids_i[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids_i == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids_i[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids_i.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = image_grid_thw[image_index]
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+ else:
+ t, h, w = video_grid_thw[video_index]
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+
+ text_len -= skip_vision_start_token[image_index - 1]
+ text_len = max(0, text_len)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+
+ if skip_vision_start_token[image_index - 1]:
+ if fix_point > 0:
+ fix_point = fix_point - st_idx
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx)
+ fix_point = 0
+ else:
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, 1, -1).expand(3, input_ids.shape[0], -1)
+ mrope_position_deltas = torch.zeros([input_ids.shape[0], 1], device=input_ids.device, dtype=input_ids.dtype)
+ return position_ids, mrope_position_deltas
+
+
+def build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config):
+ image_token_id = model_config.image_token_id
+ video_token_id = model_config.video_token_id
+ vision_start_token_id = model_config.vision_start_token_id
+ image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
+
+ boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
+ tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
+
+ messages = [{"role": "user", "content": prompt}]
+ template_caption = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + boi_token + tms_token * TIMESTEP_TOKEN_NUM
+ input_ids = tokenizer.encode(template_caption, return_tensors="pt", add_special_tokens=False)
+
+ image_grid_thw = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0)
+
+ vision_tokens = torch.zeros((1, image_len), dtype=input_ids.dtype) + image_token_id
+ vision_tokens[0, 0] = vision_start_token_id
+ input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
+
+ position_ids, _ = get_rope_index_fix_point(
+ 1,
+ image_token_id,
+ video_token_id,
+ vision_start_token_id,
+ input_ids=input_ids_pad,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=None,
+ attention_mask=None,
+ skip_vision_start_token=[1],
+ )
+
+ txt_seq_len = input_ids.shape[-1]
+ all_seq_len = position_ids.shape[-1]
+
+ token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
+ bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
+ token_types[0, bgn : bgn + image_len + TIMESTEP_TOKEN_NUM] = 1
+ token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3
+
+ vinput_mask = token_types == 1
+ token_types_bin = (token_types > 0).to(token_types.dtype)
+
+ return {
+ "input_ids": input_ids_pad,
+ "position_ids": position_ids,
+ "token_types": token_types_bin,
+ "vinput_mask": vinput_mask,
+ }
+
+
+__all__ = [
+ "CONDITION_IMAGE_SIZE",
+ "DEFAULT_TIMESTEPS",
+ "NOISE_SCALE",
+ "PATCH_SIZE",
+ "PREDEFINED_RESOLUTIONS",
+ "TIMESTEP_TOKEN_NUM",
+ "T_EPS",
+ "_calculate_dimensions",
+ "_ensure_special_tokens",
+ "_find_closest_resolution",
+ "_image_to_patch_tensor",
+ "_patch_tensor_to_image",
+ "_patches_to_np",
+ "_pil_to_normalized_tensor",
+ "_resize_pilimage",
+ "build_t2i_text_sample",
+ "get_rope_index_fix_point",
+]
diff --git a/pipelines/hidream/qwen3_vl_transformers.py b/pipelines/hidream/qwen3_vl_transformers.py
new file mode 100644
index 000000000..1abd7557f
--- /dev/null
+++ b/pipelines/hidream/qwen3_vl_transformers.py
@@ -0,0 +1,2062 @@
+# pylint: disable=unused-argument, protected-access
+import os
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Union
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+USE_BF16_ROPE = os.environ.get("USE_BF16_ROPE", "0")
+
+# Flash Attention import (FA3 preferred, FA2 fallback)
+flash_attn_version = os.environ.get("FA_VERSION", "auto")
+_flash_attn_func = None
+if flash_attn_version == "2":
+ from flash_attn import flash_attn_func as _flash_attn_func
+elif flash_attn_version == "3":
+ from flash_attn_interface import flash_attn_func as _flash_attn_func
+else:
+ try:
+ from flash_attn_interface import flash_attn_func as _flash_attn_func
+ except ImportError:
+ try:
+ from flash_attn import flash_attn_func as _flash_attn_func # pylint: disable=ungrouped-imports
+ except ImportError:
+ _flash_attn_func = None
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import check_model_inputs
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig
+
+
+def _compute_default_rope_parameters(config, device=None, seq_len=None, layer_type=None):
+ rope_theta = getattr(config, "rope_theta", 10000.0)
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+
+ if hasattr(config, "rope_parameters") and config.rope_parameters is not None:
+ if hasattr(config, "standardize_rope_params"):
+ config.standardize_rope_params()
+ rope_parameters = config.rope_parameters[layer_type] if layer_type is not None else config.rope_parameters
+ rope_theta = rope_parameters.get("rope_theta", rope_theta)
+ partial_rotary_factor = rope_parameters.get("partial_rotary_factor", partial_rotary_factor)
+
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+ inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
+ return inv_freq, 1.0
+
+class Qwen3VLVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3VLVisionPatchEmbed(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+
+ kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+ self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen3VLVisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ freqs = torch.outer(seq, self.inv_freq)
+ return freqs
+
+
+class Qwen3VLVisionPatchMerger(nn.Module):
+ def __init__(self, config: Qwen3VLVisionConfig, use_postshuffle_norm=False) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
+ self.act_fn = nn.GELU()
+ self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
+ x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+ return x
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class Qwen3VLVisionAttention(nn.Module):
+ def __init__(self, config: Qwen3VLVisionConfig) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ if self.config._attn_implementation == "flash_attention_2":
+ # Flash Attention 2: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen3VLVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.attn = Qwen3VLVisionAttention(config=config)
+ self.mlp = Qwen3VLVisionMLP(config=config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen3VLTextRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen3VLTextConfig, device=None):
+ super().__init__()
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", "default")
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ rope_init_functions = dict(ROPE_INIT_FUNCTIONS)
+ rope_init_functions.setdefault("default", _compute_default_rope_parameters)
+ self.rope_init_fn = rope_init_functions.get(self.rope_type, _compute_default_rope_parameters)
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+ # self.original_inv_freq, _ = self.rope_init_fn(self.config, device)
+
+ rope_scaling = config.rope_scaling if getattr(config, "rope_scaling", None) is not None else {}
+ self.mrope_section = rope_scaling.get("mrope_section", [24, 20, 20])
+
+ def compute_default_rope_parameters(self, config=None, device=None, seq_len=None, **kwargs):
+ rope_config = config or self.config
+ return _compute_default_rope_parameters(rope_config, device=device, seq_len=seq_len)
+
+ def _materialize_inv_freq(self, device, use_original=False):
+ source = self.original_inv_freq if use_original else self.inv_freq
+ if source is not None and not source.is_meta:
+ return source.to(device=device)
+
+ inv_freq, attention_scaling = self.rope_init_fn(self.config, device)
+ self.attention_scaling = attention_scaling
+ self.inv_freq = inv_freq
+ self.original_inv_freq = inv_freq.detach().clone()
+ return self.original_inv_freq if use_original else self.inv_freq
+
+ def apply_interleaved_mrope(self, freqs, mrope_section):
+ """Apply interleaved MRoPE to 3D rotary embeddings.
+ Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
+ interleaved [THTHWHTHW...TT], preserving frequency continuity.
+ args:
+ x: (3, bs, seq_len, head_dim // 2)
+ mrope_section: (3,)
+ returns:
+ x_t: (bs, seq_len, head_dim // 2)
+ """
+ freqs_t = freqs[0] # just overwrite the first dimension T
+ for dim, offset in enumerate((1, 2), start=1): # H, W
+ length = mrope_section[dim] * 3
+ idx = slice(offset, length, 3)
+ freqs_t[..., idx] = freqs[dim, ..., idx]
+ return freqs_t
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, Qwen3VL has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ if USE_BF16_ROPE == "1":
+ inv_freq = self._materialize_inv_freq(x.device, use_original=False)
+ inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ else:
+ inv_freq = self._materialize_inv_freq(x.device, use_original=True)
+ inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ # inv_freq_expanded = self.inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3VLTextRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Qwen3VLTextRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`, *optional*):
+ Deprecated and unused.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class Qwen3VLTextAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
+ self.k_norm = Qwen3VLTextRMSNorm(
+ self.head_dim, eps=config.rms_norm_eps
+ ) # thus post q_norm does not need reshape
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class Qwen3VLTextMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class Qwen3VLTextDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Qwen3VLTextMLP(config)
+ self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Llava outputs, with hidden states and attentions.
+ """
+)
+class Qwen3VLModelOutputWithPast(ModelOutput):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+ x_pred: Optional[torch.FloatTensor] = None
+ mid_results: Optional[list] = None
+
+
+@auto_docstring
+class Qwen3VLPreTrainedModel(PreTrainedModel):
+ config: Qwen3VLConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Qwen3VLTextDecoderLayer,
+ "attentions": Qwen3VLTextAttention,
+ }
+
+
+class Qwen3VLVisionModel(Qwen3VLPreTrainedModel):
+ config: Qwen3VLVisionConfig
+ _no_split_modules = ["Qwen3VLVisionBlock"]
+
+ def __init__(self, config, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen3VLVisionPatchEmbed(
+ config=config,
+ )
+
+ self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
+ self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen3VLVisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList([Qwen3VLVisionBlock(config) for _ in range(config.depth)])
+ self.merger = Qwen3VLVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=False,
+ )
+
+ self.deepstack_visual_indexes = config.deepstack_visual_indexes
+ self.deepstack_merger_list = nn.ModuleList(
+ [
+ Qwen3VLVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=True,
+ )
+ for _ in range(len(config.deepstack_visual_indexes))
+ ]
+ )
+
+ self.gradient_checkpointing = False
+
+ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+ merge_size = self.spatial_merge_size
+
+ max_hw = int(grid_thw[:, 1:].max().item())
+ freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2)
+ device = freq_table.device
+
+ total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
+ pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)
+
+ offset = 0
+ for num_frames, height, width in grid_thw:
+ merged_h, merged_w = height // merge_size, width // merge_size
+
+ block_rows = torch.arange(merged_h, device=device) # block row indices
+ block_cols = torch.arange(merged_w, device=device) # block col indices
+ intra_row = torch.arange(merge_size, device=device) # intra-block row offsets
+ intra_col = torch.arange(merge_size, device=device) # intra-block col offsets
+
+ # Compute full-resolution positions
+ row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
+ col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]
+
+ row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+ col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+
+ coords = torch.stack((row_idx, col_idx), dim=-1)
+
+ if num_frames > 1:
+ coords = coords.repeat(num_frames, 1)
+
+ num_tokens = coords.shape[0]
+ pos_ids[offset : offset + num_tokens] = coords
+ offset += num_tokens
+
+ embeddings = freq_table[pos_ids] # lookup rotary embeddings
+ embeddings = embeddings.flatten(1)
+ return embeddings
+
+ def fast_pos_embed_interpolate(self, grid_thw):
+ grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
+
+ idx_list = [[] for _ in range(4)]
+ weight_list = [[] for _ in range(4)]
+
+ for _t, h, w in zip(grid_ts, grid_hs, grid_ws):
+ h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
+ w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
+
+ h_idxs_floor = h_idxs.int()
+ w_idxs_floor = w_idxs.int()
+ h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+ w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+
+ dh = h_idxs - h_idxs_floor
+ dw = w_idxs - w_idxs_floor
+
+ base_h = h_idxs_floor * self.num_grid_per_side
+ base_h_ceil = h_idxs_ceil * self.num_grid_per_side
+
+ indices = [
+ (base_h[None].T + w_idxs_floor[None]).flatten(),
+ (base_h[None].T + w_idxs_ceil[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
+ ]
+
+ weights = [
+ ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
+ ((1 - dh)[None].T * dw[None]).flatten(),
+ (dh[None].T * (1 - dw)[None]).flatten(),
+ (dh[None].T * dw[None]).flatten(),
+ ]
+
+ for i in range(4):
+ idx_list[i].extend(indices[i].tolist())
+ weight_list[i].extend(weights[i].tolist())
+
+ idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
+ weight_tensor = torch.tensor(
+ weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
+ )
+ pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
+ patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
+
+ patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])
+
+ patch_pos_embeds_permute = []
+ merge_size = self.config.spatial_merge_size
+ for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
+ pos_embed = pos_embed.repeat(t, 1)
+ pos_embed = (
+ pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
+ .permute(0, 1, 3, 2, 4, 5)
+ .flatten(0, 4)
+ )
+ patch_pos_embeds_permute.append(pos_embed)
+ patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
+ return patch_pos_embeds
+
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+
+ pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+ hidden_states = hidden_states + pos_embeds
+
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ deepstack_feature_lists = []
+ for layer_num, blk in enumerate(self.blocks):
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ if layer_num in self.deepstack_visual_indexes:
+ deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
+ hidden_states
+ )
+ deepstack_feature_lists.append(deepstack_feature)
+
+ hidden_states = self.merger(hidden_states)
+
+ return hidden_states, deepstack_feature_lists
+
+
+@auto_docstring(
+ custom_intro=(
+ "Text part of Qwen3VL, "
+ "not a pure text-only model, as DeepStack integrates visual features into the early hidden states."
+ )
+)
+class Qwen3VLTextModel(Qwen3VLPreTrainedModel):
+ config: Qwen3VLTextConfig
+ _no_split_modules = ["Qwen3VLTextDecoderLayer"]
+
+ def __init__(self, config: Qwen3VLTextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @check_model_inputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ # args for deepstack
+ visual_pos_masks: Optional[torch.Tensor] = None,
+ deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
+ return_mid_results_layers: Optional[list] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Union[tuple, BaseModelOutputWithPast]:
+ r"""
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Cache positions used for KV-cache aware decoding and causal mask construction.
+ visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+ The mask of the visual positions.
+ deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+ The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+ The feature is extracted from the different visual encoder layers, and fed to the decoder
+ hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+ return_mid_results_layers (`list`, *optional*):
+ Decoder layer indices whose hidden states should be collected and returned in `mid_results`.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ text_position_ids = position_ids[0]
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ past_key_values=past_key_values,
+ position_ids=text_position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ mid_results = [] if return_mid_results_layers else None
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # === Memory profiling for decoder loop (gated by DEBUG_MEM=1) ===
+ _gc_count = 0
+
+ # decoder layers
+ for layer_idx, decoder_layer in enumerate(self.layers):
+ if self.gradient_checkpointing and torch.is_grad_enabled():
+ # Use HuggingFace's _gradient_checkpointing_func which already has
+ # use_reentrant=False baked in from gradient_checkpointing_enable().
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ )
+ _gc_count += 1
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ # add visual features to the hidden states of first several layers
+ if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+ hidden_states = self._deepstack_process(
+ hidden_states,
+ visual_pos_masks,
+ deepstack_visual_embeds[layer_idx],
+ )
+
+ if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
+ mid_results.append(hidden_states)
+
+ _a = torch.cuda.memory_allocated() / 1e9
+
+ hidden_states = self.norm(hidden_states)
+
+ output = BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+ output.mid_results = mid_results
+ return output
+
+ def _deepstack_process(
+ self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
+ ):
+ visual_pos_masks = visual_pos_masks.to(hidden_states.device)
+ visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
+ local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
+ hidden_states[visual_pos_masks, :] = local_this
+ return hidden_states
+
+class BottleneckPatchEmbed(nn.Module):
+ def __init__(self, config, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
+ super().__init__()
+ self.proj1 = nn.Linear(patch_size*patch_size*in_chans, pca_dim, bias=False)
+ self.proj2 = nn.Linear(pca_dim, embed_dim, bias=bias)
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ w1 = self.proj1.weight.data
+ nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
+ w2 = self.proj2.weight.data
+ nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
+ nn.init.constant_(self.proj2.bias, 0)
+
+ def forward(self, x):
+ x = self.proj2(self.proj1(x))
+ return x
+
+class FinalLayer(nn.Module):
+ def __init__(self, config, hidden_size, patch_size, out_channels):
+ super().__init__()
+ self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+ self.apply(self._init_weights)
+
+ def _init_weights(self, m):
+ if isinstance(m, nn.Linear):
+ nn.init.zeros_(m.weight)
+ if m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x, adaln_input=None):
+ x = self.linear(x)
+ return x
+
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, config, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ nn.init.normal_(self.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.mlp[2].weight, std=0.02)
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t * 1000, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+ return t_emb
+
+@auto_docstring
+class Qwen3VLModel(Qwen3VLPreTrainedModel):
+ base_model_prefix = ""
+ _checkpoint_conversion_mapping = {}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ config: Qwen3VLConfig
+ _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
+ self.language_model = Qwen3VLTextModel._from_config(config.text_config)
+ self.rope_deltas = None # cache rope_deltas here
+
+ self.patch_size = 32
+ self.in_channels = 3
+ hidden_size = config.text_config.hidden_size
+ bottleneck_dim = hidden_size // 4
+
+ self.t_embedder1 = TimestepEmbedder(self.config, hidden_size)
+ self.x_embedder = BottleneckPatchEmbed(self.config, patch_size = self.patch_size, in_chans = self.in_channels, pca_dim = bottleneck_dim, embed_dim = hidden_size, bias=True)
+
+ # self.t_embedder2 = TimestepEmbedder(self.config, hidden_size)
+ self.t_embedder2 = None
+ self.final_layer2 = FinalLayer(self.config, hidden_size = hidden_size, patch_size = self.patch_size, out_channels = self.in_channels)
+ self.tms_token_id = 151673
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def set_decoder(self, decoder):
+ self.language_model = decoder
+
+ def get_decoder(self):
+ return self.language_model
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids."""
+
+ # Since we use timestamps to seperate videos, like , the video_grid_thw should also be split
+ if video_grid_thw is not None:
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+ video_grid_thw[:, 0] = 1
+
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids in enumerate(total_input_ids):
+ input_ids = input_ids[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = (
+ image_grid_thw[image_index][0],
+ image_grid_thw[image_index][1],
+ image_grid_thw[image_index][2],
+ )
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+
+ else:
+ t, h, w = (
+ video_grid_thw[video_index][0],
+ video_grid_thw[video_index][1],
+ video_grid_thw[video_index][2],
+ )
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos)
+ t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+ else:
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = (
+ torch.arange(input_ids.shape[1], device=input_ids.device)
+ .view(1, 1, -1)
+ .expand(3, input_ids.shape[0], -1)
+ )
+ mrope_position_deltas = torch.zeros(
+ [input_ids.shape[0], 1],
+ device=input_ids.device,
+ dtype=input_ids.dtype,
+ )
+
+ return position_ids, mrope_position_deltas
+
+ def get_video_features(
+ self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
+ ):
+ """
+ Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned.
+
+ Args:
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ # Same implementation as for images
+ return self.get_image_features(pixel_values_videos, video_grid_thw)
+
+ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+ """
+ Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(image_embeds, split_sizes)
+ return image_embeds, deepstack_image_embeds
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: Optional[torch.FloatTensor] = None,
+ video_features: Optional[torch.FloatTensor] = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel():
+ raise ValueError(
+ f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
+ )
+
+ return special_image_mask, special_video_mask
+
+ def _run_decoder_flash(self, inputs_embeds, position_ids, token_types, return_mid_results_layers=None):
+ """Run decoder layers with flash attention two-pass approach.
+
+ Replicates the Megatron attention pattern:
+ 1. Causal attention on AR tokens only (text)
+ 2. Full (bidirectional) attention on ALL tokens
+ 3. Replace AR positions with causal result (index_copy)
+
+ This ensures AR tokens only attend causally to other AR tokens,
+ while gen tokens attend bidirectionally to everything.
+
+ Args:
+ inputs_embeds: [batch, total_seq_len, hidden]
+ position_ids: [3, batch, total_seq_len] - 3D RoPE positions
+ token_types: [batch, total_seq_len] - 0=AR, 1=gen
+ """
+ assert _flash_attn_func is not None, (
+ "Flash attention is not available. Install flash_attn_interface (FA3) or flash_attn (FA2).")
+
+ text_model = self.language_model
+
+ # Compute rotary position embeddings
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ elif position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ position_ids = position_ids[1:] # drop text_position_ids dim
+ position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids)
+ cos, sin = position_embeddings
+
+ # Precompute AR token indices (same layout for all batch items)
+ is_gen = token_types[0].bool() # [seq_len]
+ idx_ar = torch.nonzero(~is_gen, as_tuple=False).squeeze(-1) # [n_ar]
+
+ hidden_states = inputs_embeds
+ mid_results = [] if return_mid_results_layers else None
+
+ use_gc = text_model.gradient_checkpointing and torch.is_grad_enabled()
+
+ def _flash_layer_forward(hidden_states, decoder_layer, cos, sin, idx_ar):
+ """Flash attention layer forward compatible with FSDP2.
+
+ Calls decoder_layer(...) through its __call__ to trigger FSDP hooks
+ (which swap DTensor parameters to plain tensors), with self_attn.forward
+ temporarily replaced by a custom two-pass flash attention implementation.
+ """
+ original_attn_forward = decoder_layer.self_attn.forward
+
+ def _custom_flash_attn(hidden_states, position_embeddings, attention_mask=None, **kwargs):
+ attn = decoder_layer.self_attn
+ input_shape = hidden_states.shape[:-1]
+ head_dim = attn.head_dim
+ hidden_shape = (*input_shape, -1, head_dim)
+
+ # Q, K, V projections
+ q = attn.q_norm(attn.q_proj(hidden_states).view(hidden_shape))
+ k = attn.k_norm(attn.k_proj(hidden_states).view(hidden_shape))
+ v = attn.v_proj(hidden_states).view(hidden_shape)
+
+ # Apply rotary position embedding (expects [B, H, S, D])
+ cos_pe, sin_pe = position_embeddings
+ q_r = q.transpose(1, 2) # [B, H, S, D]
+ k_r = k.transpose(1, 2) # [B, KVH, S, D]
+ q_r, k_r = apply_rotary_pos_emb(q_r, k_r, cos_pe, sin_pe)
+ q = q_r.transpose(1, 2).contiguous() # [B, S, H, D]
+ k = k_r.transpose(1, 2).contiguous() # [B, S, KVH, D]
+ v = v.contiguous()
+
+ softmax_scale = head_dim ** -0.5
+
+ # --- Two-pass flash attention ---
+ # Pass 1: causal attention on AR tokens only
+ q_ar = q[:, idx_ar].contiguous()
+ k_ar = k[:, idx_ar].contiguous()
+ v_ar = v[:, idx_ar].contiguous()
+ result_ar = _flash_attn_func(q_ar.to(torch.bfloat16), k_ar.to(torch.bfloat16), v_ar.to(torch.bfloat16), softmax_scale=softmax_scale, causal=True)
+ out_ar = result_ar[0] if isinstance(result_ar, tuple) else result_ar
+
+ # Pass 2: full (bidirectional) attention on all tokens
+ result_full = _flash_attn_func(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16), softmax_scale=softmax_scale, causal=False)
+ out_full = result_full[0] if isinstance(result_full, tuple) else result_full
+
+ # Replace AR positions with causal result
+ out_full = out_full.clone()
+ out_full[:, idx_ar] = out_ar
+
+ # Output projection
+ attn_output = out_full.reshape(*input_shape, -1).contiguous()
+ attn_output = attn.o_proj(attn_output)
+ return attn_output, None
+
+ # Temporarily disable gradient checkpointing on the decoder layer
+ # to avoid nested checkpointing (the outer loop handles GC).
+ _saved_gc = decoder_layer.gradient_checkpointing
+ decoder_layer.gradient_checkpointing = False
+ decoder_layer.self_attn.forward = _custom_flash_attn
+ try:
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=(cos, sin),
+ )
+ finally:
+ decoder_layer.self_attn.forward = original_attn_forward
+ decoder_layer.gradient_checkpointing = _saved_gc
+
+ return hidden_states
+
+ for layer_idx, decoder_layer in enumerate(text_model.layers):
+ if use_gc:
+ hidden_states = torch.utils.checkpoint.checkpoint(
+ _flash_layer_forward,
+ hidden_states, decoder_layer, cos, sin, idx_ar,
+ use_reentrant=False,
+ )
+ else:
+ hidden_states = _flash_layer_forward(
+ hidden_states, decoder_layer, cos, sin, idx_ar,
+ )
+
+ if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
+ mid_results.append(hidden_states)
+
+ # Final layer norm
+ hidden_states = text_model.norm(hidden_states)
+ return hidden_states, mid_results
+
+ def _forward_generation(self, input_ids, position_ids, vinputs, timestep, token_types,
+ attention_mask=None, pixel_values=None, pixel_values_videos=None,
+ image_grid_thw=None, video_grid_thw=None, use_flash_attn=False,
+ return_mid_results_layers=None,
+ **kwargs):
+ """Forward pass for image generation (denoising step).
+
+ Args:
+ input_ids: [batch, txt_seq_len] - text token IDs (without image tokens)
+ position_ids: [3, batch, total_seq_len] - 3D RoPE positions covering text+image
+ vinputs: [batch, img_tokens, patch_dim] - patchified noise input
+ timestep: [batch] - timestep values (scalar per sample)
+ token_types: [batch, total_seq_len] or [total_seq_len, 1] - 0=AR, >0=gen
+ attention_mask: ignored (created internally for non-flash path)
+ pixel_values: optional image pixel values for conditioned generation
+ pixel_values_videos: optional video pixel values
+ image_grid_thw: optional image grid info
+ video_grid_thw: optional video grid info
+ use_flash_attn: if True, use flash attention with two-pass approach
+
+ Returns:
+ Qwen3VLModelOutputWithPast with x_pred field set.
+ """
+ # 1. Get text token embeddings
+ inputs_embeds = self.get_input_embeddings()(input_ids) # [batch, txt_seq_len, hidden]
+ batch_size, total_seq_len, _ = inputs_embeds.shape
+
+ # Parse token_types early so we can distinguish conditioning vs generation placeholders.
+ parsed_token_types = None
+ if token_types is not None:
+ if isinstance(token_types, list):
+ token_types = torch.cat(token_types, dim=0)
+ token_types = token_types.to(inputs_embeds.device)
+ if token_types.dim() == 1:
+ token_types = token_types.unsqueeze(0)
+ elif token_types.dim() == 2 and token_types.shape[-1] == 1 and token_types.shape[0] == total_seq_len:
+ # [total_seq_len, 1] -> [1, total_seq_len]
+ token_types = token_types.squeeze(-1).unsqueeze(0)
+ if token_types.shape[0] == 1 and batch_size > 1:
+ token_types = token_types.expand(batch_size, -1)
+ parsed_token_types = token_types
+
+ # 2. Process image/video embeddings if present (for image-conditioned generation)
+ if pixel_values is not None:
+ image_embeds, _ = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ if parsed_token_types is not None:
+ conditioning_image_positions = (input_ids == self.config.image_token_id) & (~parsed_token_types.bool())
+ conditioning_image_mask = conditioning_image_positions.unsqueeze(-1).expand_as(inputs_embeds)
+ if inputs_embeds[conditioning_image_mask].numel() != image_embeds.numel():
+ raise ValueError(
+ "Image features and conditioning image tokens do not match: "
+ f"tokens: {conditioning_image_positions.sum()}, features {image_embeds.shape[0]}"
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(conditioning_image_mask, image_embeds)
+ else:
+ image_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+ elif torch.is_grad_enabled():
+ # t2i task: no pixel_values, but we must run the vision encoder with a
+ # tiny dummy input so that EVERY rank has non-None (zero) gradients for
+ # vision-encoder parameters. This keeps the FSDP reduce-scatter and the
+ # replicate-group all-reduce symmetric across t2i and ref-task ranks,
+ # preventing collective hangs at backward / clip_grad_norm_.
+ # The dummy output is zeroed out before being added to inputs_embeds, so
+ # the forward result is numerically identical to the no-pixel_values path.
+ pe = self.visual.patch_embed # PatchEmbed
+ t_sz = pe.temporal_patch_size # e.g. 2
+ m_sz = self.visual.spatial_merge_size # e.g. 2
+ n_patches = t_sz * m_sz * m_sz
+ patch_dim = pe.in_channels * t_sz * pe.patch_size * pe.patch_size
+ fake_pv = torch.zeros(n_patches, patch_dim,
+ device=inputs_embeds.device,
+ dtype=pe.proj.weight.dtype)
+ fake_grid = torch.tensor([[t_sz, m_sz, m_sz]],
+ dtype=torch.long, device=inputs_embeds.device)
+ fake_embs, _ = self.get_image_features(fake_pv, fake_grid)
+ fake_embs = torch.cat(fake_embs, dim=0).to(inputs_embeds.dtype)
+ # Multiply by a zero tensor (same dtype/device) so the gradient path
+ # through the vision encoder is live but the numerical contribution is 0.
+ inputs_embeds = inputs_embeds + fake_embs.sum() * inputs_embeds.new_zeros([])
+
+ if pixel_values_videos is not None:
+ video_embeds, _ = self.get_video_features(pixel_values_videos, video_grid_thw)
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ if parsed_token_types is not None:
+ conditioning_video_positions = (input_ids == self.config.video_token_id) & (~parsed_token_types.bool())
+ conditioning_video_mask = conditioning_video_positions.unsqueeze(-1).expand_as(inputs_embeds)
+ if inputs_embeds[conditioning_video_mask].numel() != video_embeds.numel():
+ raise ValueError(
+ "Video features and conditioning video tokens do not match: "
+ f"tokens: {conditioning_video_positions.sum()}, features {video_embeds.shape[0]}"
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(conditioning_video_mask, video_embeds)
+ else:
+ _, video_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ # 3. Embed timestep and replace tms_token positions
+ if isinstance(timestep, list):
+ timestep = torch.cat(timestep, dim=0)
+ timestep = timestep.to(inputs_embeds.device)
+ t_emb = self.t_embedder1(timestep) # [batch, hidden]
+
+ tms_mask = input_ids == self.tms_token_id # [batch, txt_seq_len]
+ tms_mask_3d = tms_mask.unsqueeze(-1).expand_as(inputs_embeds)
+ t_emb_expanded = t_emb.unsqueeze(1).expand_as(inputs_embeds)
+ inputs_embeds = torch.where(tms_mask_3d, t_emb_expanded, inputs_embeds)
+
+ # 4. Embed vinputs and fill the existing multimodal placeholder slots
+ if isinstance(vinputs, list):
+ vinputs = torch.cat(vinputs, dim=0)
+ vinputs = vinputs.to(inputs_embeds.device)
+ vinputs_embedded = self.x_embedder(vinputs).to(inputs_embeds.dtype) # [batch, img_tokens, hidden]
+
+ # 5. Parse token_types to [batch, total_seq_len]
+ token_types = parsed_token_types
+ if token_types is None:
+ raise ValueError("token_types is required for generation path")
+
+ if vinputs_embedded.shape[0] == 1 and batch_size > 1:
+ vinputs_embedded = vinputs_embedded.expand(batch_size, -1, -1)
+
+ vinput_positions = token_types.bool() & ~tms_mask
+ expected_vinput_tokens = vinputs_embedded.shape[1]
+ actual_vinput_tokens = vinput_positions.sum(dim=1)
+ if not torch.all(actual_vinput_tokens == expected_vinput_tokens):
+ raise ValueError(
+ f"Vinput token count mismatch: expected {expected_vinput_tokens} placeholder tokens, got {actual_vinput_tokens.tolist()}"
+ )
+
+ for batch_index in range(batch_size):
+ inputs_embeds[batch_index, vinput_positions[batch_index]] = vinputs_embedded[batch_index]
+
+ # 6. Forward through decoder
+ mid_results = None
+
+ if use_flash_attn:
+ # Flash attention: two-pass approach (causal on AR + full on all → index_copy)
+ hidden_states, mid_results = self._run_decoder_flash(
+ inputs_embeds, position_ids, token_types,
+ return_mid_results_layers=return_mid_results_layers)
+ else:
+ # Standard path: 4D attention mask (causal for AR, full for gen tokens)
+ dtype = inputs_embeds.dtype
+ min_val = torch.finfo(dtype).min
+ attn_masks = []
+ for b in range(batch_size):
+ causal = torch.full(
+ (total_seq_len, total_seq_len), min_val,
+ device=inputs_embeds.device, dtype=dtype)
+ causal = torch.triu(causal, diagonal=1) # lower tri + diag = 0 (allowed)
+ gen_positions = token_types[b].bool() # [total_seq_len]
+ causal[gen_positions, :] = 0 # gen tokens attend to everything
+ attn_masks.append(causal)
+ attention_mask_4d = torch.stack(attn_masks, dim=0).unsqueeze(1) # [batch, 1, seq, seq]
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask_4d,
+ inputs_embeds=inputs_embeds,
+ use_cache=False,
+ return_mid_results_layers=return_mid_results_layers,
+ )
+ hidden_states = outputs.last_hidden_state
+ if hasattr(outputs, 'mid_results'):
+ mid_results = outputs.mid_results
+
+ # 7. Apply final layer to get pixel predictions
+ x_pred = self.final_layer2(hidden_states) # [batch, total_seq_len, out_dim]
+
+ return Qwen3VLModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ x_pred=x_pred,
+ mid_results=mid_results,
+ )
+
+ @auto_docstring
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ vinputs: Optional[torch.Tensor] = None,
+ timestep: Optional[torch.Tensor] = None,
+ token_types: Optional[torch.Tensor] = None,
+ use_flash_attn: bool = False,
+ return_mid_results_layers: Optional[list] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Qwen3VLModelOutputWithPast]:
+ r"""
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Cache positions used during prefilling and incremental decoding.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ vinputs (`torch.Tensor`, *optional*):
+ Patchified visual generation inputs appended to the token stream during denoising.
+ timestep (`torch.Tensor`, *optional*):
+ Per-sample denoising timesteps used to build timestep embeddings for generation.
+ token_types (`torch.Tensor`, *optional*):
+ Token-type mask distinguishing autoregressive text tokens from generation tokens.
+ use_flash_attn (`bool`, *optional*, defaults to `False`):
+ Whether to use the custom flash-attention generation path.
+ return_mid_results_layers (`list`, *optional*):
+ Decoder layer indices whose hidden states should be collected and returned in `mid_results`.
+ """
+ # Dispatch to generation forward if vinputs is provided
+ if vinputs is not None:
+ return self._forward_generation(
+ input_ids=input_ids, position_ids=position_ids,
+ vinputs=vinputs, timestep=timestep, token_types=token_types,
+ attention_mask=attention_mask,
+ pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
+ use_flash_attn=use_flash_attn,
+ return_mid_results_layers=return_mid_results_layers,
+ **kwargs)
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ image_mask = None
+ video_mask = None
+
+ if pixel_values is not None:
+ image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw)
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ visual_pos_masks = None
+ deepstack_visual_embeds = None
+ if image_mask is not None and video_mask is not None:
+ # aggregate visual_pos_masks and deepstack_visual_embeds
+ image_mask = image_mask[..., 0]
+ video_mask = video_mask[..., 0]
+ visual_pos_masks = image_mask | video_mask
+ deepstack_visual_embeds = []
+ image_mask_joint = image_mask[visual_pos_masks]
+ video_mask_joint = video_mask[visual_pos_masks]
+ for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds): # pylint: disable=possibly-used-before-assignment
+ embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)
+ embed_joint[image_mask_joint, :] = img_embed
+ embed_joint[video_mask_joint, :] = vid_embed
+ deepstack_visual_embeds.append(embed_joint)
+ elif image_mask is not None:
+ image_mask = image_mask[..., 0]
+ visual_pos_masks = image_mask
+ deepstack_visual_embeds = deepstack_image_embeds
+ elif video_mask is not None:
+ video_mask = video_mask[..., 0]
+ visual_pos_masks = video_mask
+ deepstack_visual_embeds = deepstack_video_embeds
+
+ if position_ids is None:
+ attention_mask_tensor = (
+ attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"]
+ )
+ if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
+ attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
+ # Only apply conversion for floating point tensors (inverted masks)
+ if attention_mask_tensor.dtype.is_floating_point:
+ attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min
+ attention_mask_tensor = (1.0 - attention_mask_tensor).int()
+
+ # Calculate RoPE index once per generation in the pre-fill stage only.
+ # When compiling, we can't check tensor values thus we check only input length
+ # It is safe to assume that `length!=1` means we're in pre-fill because compiled
+ # models currently cannot do asssisted decoding
+ prefill_compiled_stage = is_torchdynamo_compiling() and (
+ (input_ids is not None and input_ids.shape[1] != 1)
+ or (inputs_embeds is not None and inputs_embeds.shape[1] != 1)
+ )
+ prefill_noncompiled_stage = not is_torchdynamo_compiling() and (
+ (cache_position is not None and cache_position[0] == 0)
+ or (past_key_values is None or past_key_values.get_seq_length() == 0)
+ )
+ if (prefill_compiled_stage or prefill_noncompiled_stage) or self.rope_deltas is None:
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask=attention_mask_tensor,
+ )
+ self.rope_deltas = rope_deltas
+ # then use the prev pre-calculated rope-deltas to get the correct position ids
+ else:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ delta = (
+ (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+ if cache_position is not None
+ else 0
+ )
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ if cache_position is not None: # otherwise `deltas` is an int `0`
+ delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ visual_pos_masks=visual_pos_masks,
+ deepstack_visual_embeds=deepstack_visual_embeds,
+ **kwargs,
+ )
+
+ return Qwen3VLModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=self.rope_deltas,
+ )
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen3VL causal language model (or autoregressive) outputs.
+ """
+)
+class Qwen3VLCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+ x_pred: Optional[torch.FloatTensor] = None
+ mid_results: Optional[list] = None
+
+class HiDreamO1Qwen3VLTransformer(Qwen3VLPreTrainedModel, GenerationMixin):
+ _checkpoint_conversion_mapping = {}
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ config: Qwen3VLConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Qwen3VLModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def set_decoder(self, decoder):
+ self.model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.model.get_decoder()
+
+ def get_video_features(
+ self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
+ ):
+ return self.model.get_video_features(pixel_values_videos, video_grid_thw)
+
+ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+ return self.model.get_image_features(pixel_values, image_grid_thw)
+
+ # Make modules available through conditional class for BC
+ @property
+ def language_model(self):
+ return self.model.language_model
+
+ @property
+ def visual(self):
+ return self.model.visual
+
+ @check_model_inputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ vinputs: Optional[torch.Tensor] = None,
+ timestep: Optional[torch.Tensor] = None,
+ token_types: Optional[torch.Tensor] = None,
+ use_flash_attn: bool = False,
+ return_mid_results_layers: Optional[list] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Union[tuple, Qwen3VLCausalLMOutputWithPast]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ vinputs=vinputs,
+ timestep=timestep,
+ token_types=token_types,
+ use_flash_attn=use_flash_attn,
+ return_mid_results_layers=return_mid_results_layers,
+ **kwargs,
+ )
+
+ # Generation path: return x_pred directly
+ if vinputs is not None:
+ return Qwen3VLCausalLMOutputWithPast(
+ x_pred=outputs.x_pred,
+ mid_results=outputs.mid_results if hasattr(outputs, 'mid_results') else None,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ return Qwen3VLCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=outputs.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ cache_position=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ # Qwen3VL position_ids are prepareed with rope_deltas in forward
+ model_inputs["position_ids"] = None
+
+ if cache_position[0] != 0:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+ def _get_image_nums_and_video_nums(
+ self,
+ input_ids: Optional[torch.LongTensor],
+ inputs_embeds: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+ These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Returns:
+ image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+ video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+ """
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+
+ if inputs_embeds is not None:
+ vision_start_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ image_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ video_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ else:
+ vision_start_mask = input_ids == vision_start_token_id
+ image_mask = input_ids == image_token_id
+ video_mask = input_ids == video_token_id
+
+ vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
+ image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
+ video_nums = torch.sum(vision_first_mask & video_mask, dim=1)
+
+ return image_nums, video_nums
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: Optional[torch.LongTensor] = None,
+ **model_kwargs,
+ ) -> tuple[torch.LongTensor, dict[str, Any]]:
+ # Overwritten -- Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ video_grid_thw = model_kwargs.get("video_grid_thw", None)
+ image_nums, video_nums = self._get_image_nums_and_video_nums(
+ input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
+ )
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # split images into samples
+ samples = torch.split(image_grid_thw, list(image_nums))
+ # compute the sequence length of images for each sample
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # get the num of images for each sample
+ lengths = list(image_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "pixel_values_videos":
+ samples = torch.split(video_grid_thw, list(video_nums))
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "video_grid_thw":
+ lengths = list(video_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "second_per_grid_ts":
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size
+ )
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if (
+ key != "cache_position"
+ and dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
diff --git a/pipelines/hidream/scheduler_flashfloweuler.py b/pipelines/hidream/scheduler_flashfloweuler.py
new file mode 100644
index 000000000..c3fbe15b9
--- /dev/null
+++ b/pipelines/hidream/scheduler_flashfloweuler.py
@@ -0,0 +1,445 @@
+# Copyright 2024 Stability AI, Katherine Crowson 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.
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, is_scipy_available, logging
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+ import scipy.stats
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def hack_randn_like(*args, **kwargs):
+ # for dist broadcast
+ value = torch.randn_like(*args, **kwargs)
+ return value
+
+
+def hack_randn_tensor(*args, **kwargs):
+ # for dist broadcast
+ value = randn_tensor(*args, **kwargs)
+ return value
+
+
+@dataclass
+class FlashFlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
+ """
+ Output class for the scheduler's `step` function output.
+
+ Args:
+ prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+ Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+ denoising loop.
+ """
+
+ prev_sample: torch.FloatTensor
+
+
+class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+ """
+ Euler scheduler.
+
+ This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+ methods the library implements for all schedulers such as loading and saving.
+
+ Args:
+ num_train_timesteps (`int`, defaults to 1000):
+ The number of diffusion steps to train the model.
+ timestep_spacing (`str`, defaults to `"linspace"`):
+ The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+ shift (`float`, defaults to 1.0):
+ The shift value for the timestep schedule.
+ """
+
+ _compatibles = []
+ order = 1
+
+ @register_to_config
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ shift: float = 1.0,
+ use_dynamic_shifting=True,
+ base_shift: Optional[float] = 0.5,
+ max_shift: Optional[float] = 1.15,
+ base_image_seq_len: Optional[int] = 256,
+ max_image_seq_len: Optional[int] = 4096,
+ invert_sigmas: bool = False,
+ use_karras_sigmas: Optional[bool] = False,
+ use_exponential_sigmas: Optional[bool] = False,
+ use_beta_sigmas: Optional[bool] = False,
+ ):
+ if self.config.use_beta_sigmas and not is_scipy_available():
+ raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
+ if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
+ raise ValueError(
+ "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
+ )
+ timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+ timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+ sigmas = timesteps / num_train_timesteps
+ if not use_dynamic_shifting:
+ # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+ self.timesteps = sigmas * num_train_timesteps
+
+ self._step_index = None
+ self._begin_index = None
+
+ self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
+ self.sigma_min = self.sigmas[-1].item()
+ self.sigma_max = self.sigmas[0].item()
+
+ @property
+ def step_index(self):
+ """
+ The index counter for current timestep. It will increase 1 after each scheduler step.
+ """
+ return self._step_index
+
+ @property
+ def begin_index(self):
+ """
+ The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+ """
+ return self._begin_index
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+ def set_begin_index(self, begin_index: int = 0):
+ """
+ Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+ Args:
+ begin_index (`int`):
+ The begin index for the scheduler.
+ """
+ self._begin_index = begin_index
+
+ def scale_noise(
+ self,
+ sample: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ noise: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ """
+ Forward process in flow-matching
+
+ Args:
+ sample (`torch.FloatTensor`):
+ The input sample.
+ timestep (`int`, *optional*):
+ The current timestep in the diffusion chain.
+
+ Returns:
+ `torch.FloatTensor`:
+ A scaled input sample.
+ """
+ # Make sure sigmas and timesteps have the same device and dtype as original_samples
+ sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+ if sample.device.type == "mps" and torch.is_floating_point(timestep):
+ # mps does not support float64
+ schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+ timestep = timestep.to(sample.device, dtype=torch.float32)
+ else:
+ schedule_timesteps = self.timesteps.to(sample.device)
+ timestep = timestep.to(sample.device)
+
+ # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+ if self.begin_index is None:
+ step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+ elif self.step_index is not None:
+ # add_noise is called after first denoising step (for inpainting)
+ step_indices = [self.step_index] * timestep.shape[0]
+ else:
+ # add noise is called before first denoising step to create initial latent(img2img)
+ step_indices = [self.begin_index] * timestep.shape[0]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < len(sample.shape):
+ sigma = sigma.unsqueeze(-1)
+
+ sample = sigma * noise + (1.0 - sigma) * sample
+
+ return sample
+
+ def _sigma_to_t(self, sigma):
+ return sigma * self.config.num_train_timesteps
+
+ def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+ def set_timesteps(
+ self,
+ num_inference_steps: int | None = None,
+ device: Union[str, torch.device] = None,
+ sigmas: Optional[List[float]] = None,
+ mu: Optional[float] = None,
+ ):
+ """
+ Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+ Args:
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ """
+ if self.config.use_dynamic_shifting and mu is None:
+ raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
+
+ if sigmas is None:
+ timesteps = np.linspace(
+ self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+ )
+
+ sigmas = timesteps / self.config.num_train_timesteps
+ else:
+ sigmas = np.array(sigmas).astype(np.float32)
+ num_inference_steps = len(sigmas)
+ self.num_inference_steps = num_inference_steps
+
+ if self.config.use_dynamic_shifting:
+ sigmas = self.time_shift(mu, 1.0, sigmas)
+ else:
+ sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
+
+ if self.config.use_karras_sigmas:
+ sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+ elif self.config.use_exponential_sigmas:
+ sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+ elif self.config.use_beta_sigmas:
+ sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+ sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+ timesteps = sigmas * self.config.num_train_timesteps
+
+ if self.config.invert_sigmas:
+ sigmas = 1.0 - sigmas
+ timesteps = sigmas * self.config.num_train_timesteps
+ sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+ else:
+ sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+ self.timesteps = timesteps.to(device=device)
+ self.sigmas = sigmas
+ self._step_index = None
+ self._begin_index = None
+
+ def index_for_timestep(self, timestep, schedule_timesteps=None):
+ if schedule_timesteps is None:
+ schedule_timesteps = self.timesteps
+
+ indices = (schedule_timesteps == timestep).nonzero()
+
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ pos = 1 if len(indices) > 1 else 0
+
+ return indices[pos].item()
+
+ def _init_step_index(self, timestep):
+ if self.begin_index is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+ self._step_index = self.index_for_timestep(timestep)
+ else:
+ self._step_index = self._begin_index
+
+ def step(
+ self,
+ model_output: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ sample: torch.FloatTensor,
+ s_churn: float = 0.0,
+ s_tmin: float = 0.0,
+ s_tmax: float = float("inf"),
+ s_noise: float = 1.0,
+ noise_clip_std: float = 0.0,
+ generator: Optional[torch.Generator] = None,
+ return_dict: bool = True,
+ ) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+ process from the learned model outputs (most often the predicted noise).
+
+ Args:
+ model_output (`torch.FloatTensor`):
+ The direct output from learned diffusion model.
+ timestep (`float`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.FloatTensor`):
+ A current instance of a sample created by the diffusion process.
+ s_churn (`float`):
+ s_tmin (`float`):
+ s_tmax (`float`):
+ s_noise (`float`, defaults to 1.0):
+ Scaling factor for noise added to the sample.
+ generator (`torch.Generator`, *optional*):
+ A random number generator.
+ return_dict (`bool`):
+ Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+ tuple.
+
+ Returns:
+ [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+ If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+ returned, otherwise a tuple is returned where the first element is the sample tensor.
+ """
+
+ if (
+ isinstance(timestep, int)
+ or isinstance(timestep, torch.IntTensor)
+ or isinstance(timestep, torch.LongTensor)
+ ):
+ raise ValueError(
+ (
+ "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+ " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+ " one of the `scheduler.timesteps` as a timestep."
+ ),
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ # Upcast to avoid precision issues when computing prev_sample
+
+ sigma = self.sigmas[self.step_index]
+
+ # Upcast to avoid precision issues when computing prev_sample
+ sample = sample.to(torch.float32)
+
+ denoised = sample - model_output * sigma
+
+ if self.step_index < self.num_inference_steps:
+ sigma_next = self.sigmas[self.step_index + 1]
+ noise = hack_randn_tensor(
+ model_output.shape,
+ generator=generator,
+ device=model_output.device,
+ dtype=denoised.dtype,
+ )
+ if noise_clip_std > 0:
+ noise_std = noise.std().item()
+ clip_val = noise_clip_std * noise_std
+ noise = noise.clamp(min=-clip_val, max=clip_val)
+ sample = sigma_next * noise * s_noise + (1.0 - sigma_next) * denoised
+
+ self._step_index += 1
+ sample = sample.to(model_output.dtype)
+
+ if not return_dict:
+ return (sample,)
+
+ return FlashFlowMatchEulerDiscreteSchedulerOutput(prev_sample=sample)
+
+ # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+ def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+ """Constructs the noise schedule of Karras et al. (2022)."""
+
+ # Hack to make sure that other schedulers which copy this function don't break
+ # Add this logic to the other schedulers
+ if hasattr(self.config, "sigma_min"):
+ sigma_min = self.config.sigma_min
+ else:
+ sigma_min = None
+
+ if hasattr(self.config, "sigma_max"):
+ sigma_max = self.config.sigma_max
+ else:
+ sigma_max = None
+
+ sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+ sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+ rho = 7.0 # 7.0 is the value used in the paper
+ ramp = np.linspace(0, 1, num_inference_steps)
+ min_inv_rho = sigma_min ** (1 / rho)
+ max_inv_rho = sigma_max ** (1 / rho)
+ sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+ return sigmas
+
+ # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+ def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+ """Constructs an exponential noise schedule."""
+
+ # Hack to make sure that other schedulers which copy this function don't break
+ # Add this logic to the other schedulers
+ if hasattr(self.config, "sigma_min"):
+ sigma_min = self.config.sigma_min
+ else:
+ sigma_min = None
+
+ if hasattr(self.config, "sigma_max"):
+ sigma_max = self.config.sigma_max
+ else:
+ sigma_max = None
+
+ sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+ sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+ sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
+ return sigmas
+
+ # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+ def _convert_to_beta(
+ self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
+ ) -> torch.Tensor:
+ """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
+
+ # Hack to make sure that other schedulers which copy this function don't break
+ # Add this logic to the other schedulers
+ if hasattr(self.config, "sigma_min"):
+ sigma_min = self.config.sigma_min
+ else:
+ sigma_min = None
+
+ if hasattr(self.config, "sigma_max"):
+ sigma_max = self.config.sigma_max
+ else:
+ sigma_max = None
+
+ sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+ sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+ sigmas = np.array(
+ [
+ sigma_min + (ppf * (sigma_max - sigma_min))
+ for ppf in [
+ scipy.stats.beta.ppf(timestep, alpha, beta)
+ for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+ ]
+ ]
+ )
+ return sigmas
+
+ def __len__(self):
+ return self.config.num_train_timesteps
diff --git a/pipelines/hidream/scheduler_flowunipc.py b/pipelines/hidream/scheduler_flowunipc.py
new file mode 100644
index 000000000..b2542517b
--- /dev/null
+++ b/pipelines/hidream/scheduler_flowunipc.py
@@ -0,0 +1,778 @@
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+ SchedulerMixin,
+ SchedulerOutput)
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+ """
+ `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+ This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+ methods the library implements for all schedulers such as loading and saving.
+
+ Args:
+ num_train_timesteps (`int`, defaults to 1000):
+ The number of diffusion steps to train the model.
+ solver_order (`int`, default `2`):
+ The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+ due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+ unconditional sampling.
+ prediction_type (`str`, defaults to "flow_prediction"):
+ Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+ the flow of the diffusion process.
+ thresholding (`bool`, defaults to `False`):
+ Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+ as Stable Diffusion.
+ dynamic_thresholding_ratio (`float`, defaults to 0.995):
+ The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+ sample_max_value (`float`, defaults to 1.0):
+ The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+ predict_x0 (`bool`, defaults to `True`):
+ Whether to use the updating algorithm on the predicted x0.
+ solver_type (`str`, default `bh2`):
+ Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+ otherwise.
+ lower_order_final (`bool`, default `True`):
+ Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+ stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+ disable_corrector (`list`, default `[]`):
+ Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+ and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+ usually disabled during the first few steps.
+ solver_p (`SchedulerMixin`, default `None`):
+ Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+ use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+ Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+ the sigmas are determined according to a sequence of noise levels {sigma_i}.
+ use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+ Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+ timestep_spacing (`str`, defaults to `"linspace"`):
+ The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+ steps_offset (`int`, defaults to 0):
+ An offset added to the inference steps, as required by some model families.
+ final_sigmas_type (`str`, defaults to `"zero"`):
+ The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+ sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+ """
+
+ _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+ order = 1
+
+ @register_to_config
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ solver_order: int = 2,
+ prediction_type: str = "flow_prediction",
+ shift: Optional[float] = 1.0,
+ use_dynamic_shifting=False,
+ thresholding: bool = False,
+ dynamic_thresholding_ratio: float = 0.995,
+ sample_max_value: float = 1.0,
+ predict_x0: bool = True,
+ solver_type: str = "bh2",
+ lower_order_final: bool = True,
+ disable_corrector: List[int] = [],
+ solver_p: SchedulerMixin = None,
+ timestep_spacing: str = "linspace",
+ steps_offset: int = 0,
+ final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
+ ):
+
+ if solver_type not in ["bh1", "bh2"]:
+ if solver_type in ["midpoint", "heun", "logrho"]:
+ self.register_to_config(solver_type="bh2")
+ else:
+ raise NotImplementedError(
+ f"{solver_type} is not implemented for {self.__class__}")
+
+ self.predict_x0 = predict_x0
+ # setable values
+ self.num_inference_steps = None
+ alphas = np.linspace(1, 1 / num_train_timesteps,
+ num_train_timesteps)[::-1].copy()
+ sigmas = 1.0 - alphas
+ sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+ if not use_dynamic_shifting:
+ # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ self.sigmas = sigmas
+ self.timesteps = sigmas * num_train_timesteps
+
+ self.model_outputs = [None] * solver_order
+ self.timestep_list = [None] * solver_order
+ self.lower_order_nums = 0
+ self.disable_corrector = disable_corrector
+ self.solver_p = solver_p
+ self.last_sample = None
+ self._step_index = None
+ self._begin_index = None
+
+ self.sigmas = self.sigmas.to(
+ "cpu") # to avoid too much CPU/GPU communication
+ self.sigma_min = self.sigmas[-1].item()
+ self.sigma_max = self.sigmas[0].item()
+
+ @property
+ def step_index(self):
+ """
+ The index counter for current timestep. It will increase 1 after each scheduler step.
+ """
+ return self._step_index
+
+ @property
+ def begin_index(self):
+ """
+ The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+ """
+ return self._begin_index
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+ def set_begin_index(self, begin_index: int = 0):
+ """
+ Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+ Args:
+ begin_index (`int`):
+ The begin index for the scheduler.
+ """
+ self._begin_index = begin_index
+
+ # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+ def set_timesteps(
+ self,
+ num_inference_steps: Union[int, None] = None,
+ device: Union[str, torch.device] = None,
+ sigmas: Optional[List[float]] = None,
+ mu: Optional[Union[float, None]] = None,
+ shift: Optional[Union[float, None]] = None,
+ ):
+ """
+ Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+ Args:
+ num_inference_steps (`int`):
+ Total number of the spacing of the time steps.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ """
+
+ if self.config.use_dynamic_shifting and mu is None:
+ raise ValueError(
+ " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+ )
+
+ if sigmas is None:
+ sigmas = np.linspace(self.sigma_max, self.sigma_min,
+ num_inference_steps +
+ 1).copy()[:-1] # pyright: ignore
+
+ if self.config.use_dynamic_shifting:
+ sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore
+ else:
+ if shift is None:
+ shift = self.config.shift
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ if self.config.final_sigmas_type == "sigma_min":
+ sigma_last = ((1 - self.alphas_cumprod[0]) /
+ self.alphas_cumprod[0])**0.5
+ elif self.config.final_sigmas_type == "zero":
+ sigma_last = 0
+ else:
+ raise ValueError(
+ f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+ )
+
+ timesteps = sigmas * self.config.num_train_timesteps
+ sigmas = np.concatenate([sigmas, [sigma_last]
+ ]).astype(np.float32) # pyright: ignore
+
+ self.sigmas = torch.from_numpy(sigmas)
+ self.timesteps = torch.from_numpy(timesteps).to(
+ device=device, dtype=torch.int64)
+
+ self.num_inference_steps = len(timesteps)
+
+ self.model_outputs = [
+ None,
+ ] * self.config.solver_order
+ self.lower_order_nums = 0
+ self.last_sample = None
+ if self.solver_p:
+ self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+ # add an index counter for schedulers that allow duplicated timesteps
+ self._step_index = None
+ self._begin_index = None
+ self.sigmas = self.sigmas.to(
+ "cpu") # to avoid too much CPU/GPU communication
+
+ # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+ """
+ "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+ prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+ s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+ pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+ photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+ https://arxiv.org/abs/2205.11487
+ """
+ dtype = sample.dtype
+ batch_size, channels, *remaining_dims = sample.shape
+
+ if dtype not in (torch.float32, torch.float64):
+ sample = sample.float(
+ ) # upcast for quantile calculation, and clamp not implemented for cpu half
+
+ # Flatten sample for doing quantile calculation along each image
+ sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+ abs_sample = sample.abs() # "a certain percentile absolute pixel value"
+
+ s = torch.quantile(
+ abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+ s = torch.clamp(
+ s, min=1, max=self.config.sample_max_value
+ ) # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+ s = s.unsqueeze(
+ 1) # (batch_size, 1) because clamp will broadcast along dim=0
+ sample = torch.clamp(
+ sample, -s, s
+ ) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+ sample = sample.reshape(batch_size, channels, *remaining_dims)
+ sample = sample.to(dtype)
+
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+ def _sigma_to_t(self, sigma):
+ return sigma * self.config.num_train_timesteps
+
+ def _sigma_to_alpha_sigma_t(self, sigma):
+ return 1 - sigma, sigma
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+ def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+ def convert_model_output(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ r"""
+ Convert the model output to the corresponding type the UniPC algorithm needs.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model.
+ timestep (`int`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+
+ Returns:
+ `torch.Tensor`:
+ The converted model output.
+ """
+ args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+ if sample is None:
+ if len(args) > 1:
+ sample = args[1]
+ else:
+ raise ValueError(
+ "missing `sample` as a required keyward argument")
+
+ sigma = self.sigmas[self.step_index]
+ _alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+ if self.predict_x0:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ x0_pred = self._threshold_sample(x0_pred)
+
+ return x0_pred
+ else:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ epsilon = sample - (1 - sigma_t) * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ x0_pred = self._threshold_sample(x0_pred)
+ epsilon = model_output + x0_pred
+
+ return epsilon
+
+ def multistep_uni_p_bh_update(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ order: int | None = None, # pyright: ignore
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model at the current timestep.
+ prev_timestep (`int`):
+ The previous discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ order (`int`):
+ The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+ Returns:
+ `torch.Tensor`:
+ The sample tensor at the previous timestep.
+ """
+ args[0] if len(args) > 0 else kwargs.pop(
+ "prev_timestep", None)
+ if sample is None:
+ if len(args) > 1:
+ sample = args[1]
+ else:
+ raise ValueError(
+ " missing `sample` as a required keyward argument")
+ if order is None:
+ if len(args) > 2:
+ order = args[2]
+ else:
+ raise ValueError(
+ " missing `order` as a required keyward argument")
+ model_output_list = self.model_outputs
+
+ s0 = self.timestep_list[-1]
+ m0 = model_output_list[-1]
+ x = sample
+
+ if self.solver_p:
+ x_t = self.solver_p.step(model_output, s0, x).prev_sample
+ return x_t
+
+ sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
+ self.step_index] # pyright: ignore
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+ h = lambda_t - lambda_s0
+ device = sample.device
+
+ rks = []
+ D1s = []
+ for i in range(1, order):
+ si = self.step_index - i # pyright: ignore
+ mi = model_output_list[-(i + 1)]
+ alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+ lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+ rk = (lambda_si - lambda_s0) / h
+ rks.append(rk)
+ D1s.append((mi - m0) / rk) # pyright: ignore
+
+ rks.append(1.0)
+ rks = torch.tensor(rks, device=device)
+
+ R = []
+ b = []
+
+ hh = -h if self.predict_x0 else h
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+ h_phi_k = h_phi_1 / hh - 1
+
+ factorial_i = 1
+
+ if self.config.solver_type == "bh1":
+ B_h = hh
+ elif self.config.solver_type == "bh2":
+ B_h = torch.expm1(hh)
+ else:
+ raise NotImplementedError
+
+ for i in range(1, order + 1):
+ R.append(torch.pow(rks, i - 1))
+ b.append(h_phi_k * factorial_i / B_h)
+ factorial_i *= i + 1
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+ R = torch.stack(R)
+ b = torch.tensor(b, device=device)
+
+ if len(D1s) > 0:
+ D1s = torch.stack(D1s, dim=1) # (B, K)
+ # for order 2, we use a simplified version
+ if order == 2:
+ rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+ else:
+ rhos_p = torch.linalg.solve(R[:-1, :-1],
+ b[:-1]).to(device).to(x.dtype)
+ else:
+ D1s = None
+
+ if self.predict_x0:
+ x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+ if D1s is not None:
+ pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+ D1s) # pyright: ignore
+ else:
+ pred_res = 0
+ x_t = x_t_ - alpha_t * B_h * pred_res
+ else:
+ x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+ if D1s is not None:
+ pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+ D1s) # pyright: ignore
+ else:
+ pred_res = 0
+ x_t = x_t_ - sigma_t * B_h * pred_res
+
+ x_t = x_t.to(x.dtype)
+ return x_t
+
+ def multistep_uni_c_bh_update(
+ self,
+ this_model_output: torch.Tensor,
+ *args,
+ last_sample: torch.Tensor = None,
+ this_sample: torch.Tensor = None,
+ order: int | None = None, # pyright: ignore
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the UniC (B(h) version).
+
+ Args:
+ this_model_output (`torch.Tensor`):
+ The model outputs at `x_t`.
+ this_timestep (`int`):
+ The current timestep `t`.
+ last_sample (`torch.Tensor`):
+ The generated sample before the last predictor `x_{t-1}`.
+ this_sample (`torch.Tensor`):
+ The generated sample after the last predictor `x_{t}`.
+ order (`int`):
+ The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+ Returns:
+ `torch.Tensor`:
+ The corrected sample tensor at the current timestep.
+ """
+ args[0] if len(args) > 0 else kwargs.pop(
+ "this_timestep", None)
+ if last_sample is None:
+ if len(args) > 1:
+ last_sample = args[1]
+ else:
+ raise ValueError(
+ " missing`last_sample` as a required keyward argument")
+ if this_sample is None:
+ if len(args) > 2:
+ this_sample = args[2]
+ else:
+ raise ValueError(
+ " missing`this_sample` as a required keyward argument")
+ if order is None:
+ if len(args) > 3:
+ order = args[3]
+ else:
+ raise ValueError(
+ " missing`order` as a required keyward argument")
+
+ model_output_list = self.model_outputs
+
+ m0 = model_output_list[-1]
+ x = last_sample
+ x_t = this_sample
+ model_t = this_model_output
+
+ sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+ self.step_index - 1] # pyright: ignore
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+ h = lambda_t - lambda_s0
+ device = this_sample.device
+
+ rks = []
+ D1s = []
+ for i in range(1, order):
+ si = self.step_index - (i + 1) # pyright: ignore
+ mi = model_output_list[-(i + 1)]
+ alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+ lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+ rk = (lambda_si - lambda_s0) / h
+ rks.append(rk)
+ D1s.append((mi - m0) / rk) # pyright: ignore
+
+ rks.append(1.0)
+ rks = torch.tensor(rks, device=device)
+
+ R = []
+ b = []
+
+ hh = -h if self.predict_x0 else h
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+ h_phi_k = h_phi_1 / hh - 1
+
+ factorial_i = 1
+
+ if self.config.solver_type == "bh1":
+ B_h = hh
+ elif self.config.solver_type == "bh2":
+ B_h = torch.expm1(hh)
+ else:
+ raise NotImplementedError
+
+ for i in range(1, order + 1):
+ R.append(torch.pow(rks, i - 1))
+ b.append(h_phi_k * factorial_i / B_h)
+ factorial_i *= i + 1
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+ R = torch.stack(R)
+ b = torch.tensor(b, device=device)
+
+ if len(D1s) > 0:
+ D1s = torch.stack(D1s, dim=1)
+ else:
+ D1s = None
+
+ # for order 1, we use a simplified version
+ if order == 1:
+ rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+ else:
+ rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+ if self.predict_x0:
+ x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+ if D1s is not None:
+ corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+ else:
+ corr_res = 0
+ D1_t = model_t - m0
+ x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+ else:
+ x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+ if D1s is not None:
+ corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+ else:
+ corr_res = 0
+ D1_t = model_t - m0
+ x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+ x_t = x_t.to(x.dtype)
+ return x_t
+
+ def index_for_timestep(self, timestep, schedule_timesteps=None):
+ if schedule_timesteps is None:
+ schedule_timesteps = self.timesteps
+
+ indices = (schedule_timesteps == timestep).nonzero()
+
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ pos = 1 if len(indices) > 1 else 0
+
+ return indices[pos].item()
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+ def _init_step_index(self, timestep):
+ """
+ Initialize the step_index counter for the scheduler.
+ """
+
+ if self.begin_index is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+ self._step_index = self.index_for_timestep(timestep)
+ else:
+ self._step_index = self._begin_index
+
+ def step(self,
+ model_output: torch.Tensor,
+ timestep: Union[int, torch.Tensor],
+ sample: torch.Tensor,
+ return_dict: bool = True,
+ generator=None) -> Union[SchedulerOutput, Tuple]:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+ the multistep UniPC.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from learned diffusion model.
+ timestep (`int`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ return_dict (`bool`):
+ Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+ Returns:
+ [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+ If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+ tuple is returned where the first element is the sample tensor.
+
+ """
+ if self.num_inference_steps is None:
+ raise ValueError(
+ "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ use_corrector = (
+ self.step_index > 0 and
+ self.step_index - 1 not in self.disable_corrector and
+ self.last_sample is not None # pyright: ignore
+ )
+
+ model_output_convert = self.convert_model_output(
+ model_output, sample=sample)
+ if use_corrector:
+ sample = self.multistep_uni_c_bh_update(
+ this_model_output=model_output_convert,
+ last_sample=self.last_sample,
+ this_sample=sample,
+ order=self.this_order,
+ )
+
+ for i in range(self.config.solver_order - 1):
+ self.model_outputs[i] = self.model_outputs[i + 1]
+ self.timestep_list[i] = self.timestep_list[i + 1]
+
+ self.model_outputs[-1] = model_output_convert
+ self.timestep_list[-1] = timestep # pyright: ignore
+
+ if self.config.lower_order_final:
+ this_order = min(self.config.solver_order,
+ len(self.timesteps) -
+ self.step_index) # pyright: ignore
+ else:
+ this_order = self.config.solver_order
+
+ self.this_order = min(this_order,
+ self.lower_order_nums + 1) # warmup for multistep
+ assert self.this_order > 0
+
+ self.last_sample = sample
+ prev_sample = self.multistep_uni_p_bh_update(
+ model_output=model_output, # pass the original non-converted model output, in case solver-p is used
+ sample=sample,
+ order=self.this_order,
+ )
+
+ if self.lower_order_nums < self.config.solver_order:
+ self.lower_order_nums += 1
+
+ # upon completion increase step index by one
+ self._step_index += 1 # pyright: ignore
+
+ if not return_dict:
+ return (prev_sample,)
+
+ return SchedulerOutput(prev_sample=prev_sample)
+
+ def scale_model_input(self, sample: torch.Tensor, *args,
+ **kwargs) -> torch.Tensor:
+ """
+ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+ current timestep.
+
+ Args:
+ sample (`torch.Tensor`):
+ The input sample.
+
+ Returns:
+ `torch.Tensor`:
+ A scaled input sample.
+ """
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+ def add_noise(
+ self,
+ original_samples: torch.Tensor,
+ noise: torch.Tensor,
+ timesteps: torch.IntTensor,
+ ) -> torch.Tensor:
+ # Make sure sigmas and timesteps have the same device and dtype as original_samples
+ sigmas = self.sigmas.to(
+ device=original_samples.device, dtype=original_samples.dtype)
+ if original_samples.device.type == "mps" and torch.is_floating_point(
+ timesteps):
+ # mps does not support float64
+ schedule_timesteps = self.timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ timesteps = timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ else:
+ schedule_timesteps = self.timesteps.to(original_samples.device)
+ timesteps = timesteps.to(original_samples.device)
+
+ # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+ if self.begin_index is None:
+ step_indices = [
+ self.index_for_timestep(t, schedule_timesteps)
+ for t in timesteps
+ ]
+ elif self.step_index is not None:
+ # add_noise is called after first denoising step (for inpainting)
+ step_indices = [self.step_index] * timesteps.shape[0]
+ else:
+ # add noise is called before first denoising step to create initial latent(img2img)
+ step_indices = [self.begin_index] * timesteps.shape[0]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < len(original_samples.shape):
+ sigma = sigma.unsqueeze(-1)
+
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+ noisy_samples = alpha_t * original_samples + sigma_t * noise
+ return noisy_samples
+
+ def __len__(self):
+ return self.config.num_train_timesteps
diff --git a/pipelines/lumina_dimmo/lumina_dimoo.py b/pipelines/lumina_dimmo/lumina_dimoo.py
new file mode 100644
index 000000000..649e34d0a
--- /dev/null
+++ b/pipelines/lumina_dimmo/lumina_dimoo.py
@@ -0,0 +1,2679 @@
+# 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.
+
+import functools
+import math
+import random
+import sys
+from abc import abstractmethod
+from dataclasses import dataclass, fields
+from enum import Enum
+from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union, cast
+from accelerate import init_empty_weights
+
+from tqdm.rich import tqdm
+import numpy as np
+import torch
+import torch.backends.cuda
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+from PIL import Image, ImageDraw
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedModel
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from diffusers import DiffusionPipeline, VQModel
+from diffusers.utils import BaseOutput, logging, replace_example_docstring
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+# --- Start of model definition copied from Lumina-DiMOO ---
+
+
+class StrEnum(str, Enum):
+ """
+ This is equivalent to Python's :class:`enum.StrEnum` since version 3.11.
+ We include this here for compatibility with older version of Python.
+ """
+
+ def __str__(self) -> str:
+ return self.value
+
+ def __repr__(self) -> str:
+ return f"'{str(self)}'"
+
+
+class LayerNormType(StrEnum):
+ default = "default"
+ low_precision = "low_precision"
+ rms = "rms"
+ gemma_rms = "gemma_rms"
+ amd_compatible = "amd_compatible"
+
+
+class ActivationType(StrEnum):
+ gelu = "gelu"
+ relu = "relu"
+ silu = "silu"
+ swiglu = "swiglu"
+
+
+class BlockType(StrEnum):
+ sequential = "sequential"
+ parallel = "parallel"
+ llama = "llama"
+
+
+class InitFnType(StrEnum):
+ mitchell = "mitchell"
+ normal = "normal"
+ kaiming_normal = "kaiming_normal"
+ fan_in = "fan_in"
+ full_megatron = "full_megatron"
+
+
+@dataclass
+class ModelConfig:
+ """
+ LLaDA (model) configuration.
+ """
+
+ # Note that the defaults for these attributes are equivalent to the base GPT2 model.
+
+ d_model: int = 768
+ n_heads: int = 12
+ n_kv_heads: Optional[int] = None
+ n_layers: int = 12
+ mlp_ratio: int = 4
+ mlp_hidden_size: Optional[int] = None
+ activation_type: ActivationType = ActivationType.swiglu
+ block_type: BlockType = BlockType.sequential
+ block_group_size: int = 1
+ alibi: bool = False
+ alibi_bias_max: float = 8.0
+ rope: bool = False
+ rope_full_precision: bool = True
+ flash_attention: bool = False
+ attention_dropout: float = 0.1
+ multi_query_attention: Optional[bool] = None
+ attention_layer_norm: bool = False
+ residual_dropout: float = 0.1
+ embedding_dropout: float = 0.1
+ input_emb_norm: bool = False
+ layer_norm_type: LayerNormType = LayerNormType.default
+ layer_norm_with_affine: bool = True
+ rms_norm_eps: float = 1e-05
+ attention_layer_norm_with_affine: bool = True
+ max_sequence_length: int = 1024
+ rope_theta: float = 10000.0
+ include_qkv_bias: Optional[bool] = False
+ include_bias: bool = False
+ bias_for_layer_norm: Optional[bool] = None
+ scale_logits: bool = False
+ vocab_size: int = 50257
+ embedding_size: Optional[int] = 50304
+ weight_tying: bool = True
+ eos_token_id: int = 50256
+ pad_token_id: int = 50256
+ mask_token_id: Optional[int] = 50256
+ init_device: Optional[str] = None
+ init_fn: InitFnType = InitFnType.normal
+ init_std: float = 0.02
+ init_cutoff_factor: Optional[float] = None
+ precision: Optional[str] = None
+
+ @property
+ def effective_n_kv_heads(self) -> int:
+ if self.n_kv_heads is None:
+ if self.multi_query_attention is True:
+ return 1
+ else:
+ return self.n_heads
+ else:
+ if self.multi_query_attention is None:
+ return self.n_kv_heads
+ if self.multi_query_attention:
+ n_kv_heads_should_be = 1
+ else:
+ n_kv_heads_should_be = self.n_heads
+ if self.n_kv_heads == n_kv_heads_should_be:
+ return n_kv_heads_should_be
+ else:
+ raise Exception("You can't set `multi_query_attention` and `n_kv_heads` at the same time.")
+
+
+class ActivationCheckpointingStrategy(StrEnum):
+ whole_layer = "whole_layer"
+ one_in_two = "one_in_two"
+ one_in_three = "one_in_three"
+ one_in_four = "one_in_four"
+ two_in_three = "two_in_three"
+ three_in_four = "three_in_four"
+ four_in_five = "four_in_five"
+ nine_in_ten = "nine_in_ten"
+ fine_grained = "fine_grained"
+
+
+class LLaDAConfig(PretrainedConfig):
+ model_type = "llada"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(self, use_cache: bool = False, **kwargs):
+ model_config = ModelConfig()
+ all_kwargs = model_config.__dict__
+ all_kwargs.update(kwargs)
+ all_kwargs.update({"use_cache": use_cache})
+ all_kwargs.update({"architectures": all_kwargs.get("architectures", ["LLaDAModelLM"])})
+ super().__init__(**all_kwargs)
+
+ @property
+ def num_attention_heads(self):
+ return self.n_heads
+
+ @property
+ def num_hidden_layers(self):
+ return self.n_layers
+
+ @property
+ def hidden_size(self):
+ return self.d_model
+
+
+if sys.version_info >= (3, 9):
+ from collections.abc import MutableMapping
+elif sys.version_info >= (3, 8):
+ from typing import MutableMapping
+else:
+ raise SystemExit("This script supports Python 3.8 or higher")
+
+
+class ModuleType(StrEnum):
+ in_module = "in"
+ out_module = "out"
+ emb = "emb"
+ final_out = "final_out"
+
+
+def init_weights(
+ config: ModelConfig,
+ module: Union[nn.Linear, nn.Embedding],
+ d: Optional[int] = None,
+ layer_id: Optional[int] = None,
+ std_factor: float = 1.0,
+ type_of_module: Optional[ModuleType] = None,
+) -> None:
+ d = d if d is not None else config.d_model
+ if config.init_fn == InitFnType.normal:
+ std = config.init_std * std_factor
+ if config.init_cutoff_factor is not None:
+ cutoff_value = config.init_cutoff_factor * std
+ nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
+ else:
+ nn.init.normal_(module.weight, mean=0.0, std=std)
+ elif config.init_fn == InitFnType.mitchell:
+ std = std_factor / math.sqrt(d)
+ if layer_id is not None:
+ std = std / math.sqrt(2 * (layer_id + 1))
+ nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)
+ elif config.init_fn == InitFnType.kaiming_normal:
+ nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
+ elif config.init_fn == InitFnType.fan_in:
+ std = std_factor / math.sqrt(d)
+ nn.init.normal_(module.weight, mean=0.0, std=std)
+ elif config.init_fn == InitFnType.full_megatron:
+ if type_of_module is None:
+ raise RuntimeError(f"When using the {InitFnType.full_megatron} init, every module must have a type.")
+
+ cutoff_factor = config.init_cutoff_factor
+ if cutoff_factor is None:
+ cutoff_factor = 3
+
+ if type_of_module == ModuleType.in_module:
+ # for att_proj (same as QKV), ff_proj
+ std = config.init_std
+ elif type_of_module == ModuleType.out_module:
+ # for attn_out, ff_out
+ std = config.init_std / math.sqrt(2.0 * config.n_layers)
+ elif type_of_module == ModuleType.emb:
+ # positional embeddings (wpe)
+ # token embeddings (wte)
+ std = config.init_std
+ elif type_of_module == ModuleType.final_out:
+ # final output (ff_out)
+ std = config.d_model**-0.5
+ else:
+ raise RuntimeError(f"Unknown module type '{type_of_module}'")
+ nn.init.trunc_normal_(
+ module.weight,
+ mean=0.0,
+ std=std,
+ a=-cutoff_factor * std,
+ b=cutoff_factor * std,
+ )
+ else:
+ raise NotImplementedError(config.init_fn)
+
+ if isinstance(module, nn.Linear):
+ if module.bias is not None:
+ nn.init.zeros_(module.bias)
+
+ if config.init_fn == InitFnType.normal and getattr(module, "_is_residual", False):
+ with torch.no_grad():
+ module.weight.div_(math.sqrt(2 * config.n_layers))
+
+
+def ensure_finite_(x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False):
+ if check_neg_inf:
+ x.masked_fill_(x == float("-inf"), torch.finfo(x.dtype).min)
+ if check_pos_inf:
+ x.masked_fill_(x == float("inf"), torch.finfo(x.dtype).max)
+
+
+def activation_checkpoint_function(cfg: ModelConfig):
+ preserve_rng_state = (
+ (cfg.attention_dropout == 0.0) and (cfg.embedding_dropout == 0.0) and (cfg.residual_dropout == 0.0)
+ )
+ from torch.utils.checkpoint import checkpoint
+
+ return functools.partial(
+ checkpoint,
+ preserve_rng_state=preserve_rng_state,
+ use_reentrant=False,
+ )
+
+
+class BufferCache(dict, MutableMapping[str, torch.Tensor]):
+ """
+ Cache for attention biases and other things that would normally be stored as buffers.
+ We avoid using buffers because we've run into various issues doing so with FSDP.
+ In general it appears the way FSDP handles buffers is not well-defined.
+ It doesn't shard them but apparently it does synchronize them across processes, which we want to avoid
+ since (A) it isn't necessary, and (B) we sometimes have `-inf` in these biases which might get turned into
+ NaNs when they're synchronized due to casting or some other issue.
+ """
+
+
+def _non_meta_init_device(config: ModelConfig) -> torch.device:
+ if config.init_device is not None and config.init_device != "meta":
+ return torch.device(config.init_device)
+ else:
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Dropout(nn.Dropout):
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ if self.p == 0.0:
+ return input
+ else:
+ return F.dropout(input, self.p, self.training, self.inplace)
+
+
+class LayerNormBase(nn.Module):
+ def __init__(
+ self,
+ config: ModelConfig,
+ *,
+ size: Optional[int] = None,
+ elementwise_affine: Optional[bool] = True,
+ eps: float = 1e-05,
+ ):
+ super().__init__()
+ self.config = config
+ self.eps = eps
+ self.normalized_shape = (size or config.d_model,)
+ if elementwise_affine or (elementwise_affine is None and self.config.layer_norm_with_affine):
+ self.weight = nn.Parameter(torch.ones(self.normalized_shape, device=config.init_device))
+ use_bias = self.config.bias_for_layer_norm
+ if use_bias is None:
+ use_bias = self.config.include_bias
+ if use_bias:
+ self.bias = nn.Parameter(torch.zeros(self.normalized_shape, device=config.init_device))
+ else:
+ self.register_parameter("bias", None)
+ else:
+ self.register_parameter("bias", None)
+ self.register_parameter("weight", None)
+
+ @abstractmethod
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ raise NotImplementedError
+
+ @classmethod
+ def build(cls, config: ModelConfig, size: Optional[int] = None, **kwargs) -> "LayerNormBase":
+ if config.layer_norm_type == LayerNormType.default:
+ return LayerNorm(config, size=size, low_precision=False, **kwargs)
+ elif config.layer_norm_type == LayerNormType.low_precision:
+ return LayerNorm(config, size=size, low_precision=True, **kwargs)
+ elif config.layer_norm_type == LayerNormType.rms:
+ return RMSLayerNorm(config, size=size, **kwargs)
+ elif config.layer_norm_type == LayerNormType.gemma_rms:
+ return GemmaRMSLayerNorm(config, size=size, **kwargs)
+ else:
+ raise NotImplementedError(f"Unknown LayerNorm type: '{config.layer_norm_type}'")
+
+ def _cast_if_autocast_enabled(self, tensor: torch.Tensor, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+ if tensor.device.type == "cuda" and torch.is_autocast_enabled():
+ return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_gpu_dtype())
+ elif tensor.device.type == "cpu" and torch.is_autocast_cpu_enabled():
+ return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_cpu_dtype())
+ else:
+ return tensor
+
+ def reset_parameters(self):
+ if self.weight is not None:
+ torch.nn.init.ones_(self.weight) # type: ignore
+ if self.bias is not None:
+ torch.nn.init.zeros_(self.bias) # type: ignore
+
+
+class LayerNorm(LayerNormBase):
+ def __init__(
+ self,
+ config: ModelConfig,
+ size: Optional[int] = None,
+ low_precision: bool = False,
+ elementwise_affine: Optional[bool] = None,
+ eps: float = 1e-05,
+ ):
+ super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps)
+ self.low_precision = low_precision
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.low_precision:
+ module_device = x.device
+ downcast_x = self._cast_if_autocast_enabled(x)
+ downcast_weight = (
+ self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+ )
+ downcast_bias = self._cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+ with torch.autocast(enabled=False, device_type=module_device.type):
+ return F.layer_norm(
+ downcast_x, self.normalized_shape, weight=downcast_weight, bias=downcast_bias, eps=self.eps
+ )
+ else:
+ return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps)
+
+
+class RMSLayerNorm(LayerNormBase):
+ def __init__(
+ self,
+ config: ModelConfig,
+ size: Optional[int] = None,
+ elementwise_affine: Optional[bool] = None,
+ eps: float = 1e-5,
+ ):
+ super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=config.rms_norm_eps)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ with torch.autocast(enabled=False, device_type=x.device.type):
+ og_dtype = x.dtype
+ x = x.to(torch.float32)
+ variance = x.pow(2).mean(-1, keepdim=True)
+ x = x * torch.rsqrt(variance + self.eps)
+ x = x.to(og_dtype)
+
+ if self.weight is not None:
+ if self.bias is not None:
+ return self.weight * x + self.bias
+ else:
+ return self.weight * x
+ else:
+ return x
+
+
+class GemmaRMSLayerNorm(LayerNormBase):
+ def __init__(
+ self,
+ config: ModelConfig,
+ size: Optional[int] = None,
+ elementwise_affine: Optional[bool] = None,
+ eps: float = 1e-5,
+ ):
+ super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=config.rms_norm_eps)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ with torch.autocast(enabled=False, device_type=x.device.type):
+ og_dtype = x.dtype
+ x = x.to(torch.float32)
+ variance = x.pow(2).mean(-1, keepdim=True)
+ x = x * torch.rsqrt(variance + self.eps)
+ x = x.to(og_dtype)
+
+ if self.weight is not None:
+ if self.bias is not None:
+ return x * (1 + self.weight) + self.bias
+ else:
+ return x * (1 + self.weight)
+ else:
+ return x
+
+
+class RotaryEmbedding(nn.Module):
+ def __init__(self, config: ModelConfig, cache: BufferCache):
+ super().__init__()
+ self.config = config
+ self.__cache = cache
+ # Warm up cache.
+ self.rope_theta = config.rope_theta
+ self.get_rotary_embedding(config.max_sequence_length, _non_meta_init_device(config))
+
+ def get_rotary_embedding(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+ if (
+ (pos_sin := self.__cache.get("rope_pos_sin")) is not None
+ and (pos_cos := self.__cache.get("rope_pos_cos")) is not None
+ and pos_sin.shape[-2] >= seq_len
+ and pos_cos.shape[-2] >= seq_len
+ ):
+ if pos_sin.device != device:
+ pos_sin = pos_sin.to(device)
+ self.__cache["rope_pos_sin"] = pos_sin
+ if pos_cos.device != device:
+ pos_cos = pos_cos.to(device)
+ self.__cache["rope_pos_cos"] = pos_cos
+ return pos_sin[:, :, :seq_len, :], pos_cos[:, :, :seq_len, :]
+
+ with torch.autocast(device.type, enabled=False):
+ dim = self.config.d_model // self.config.n_heads
+ inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float) / dim))
+ seq = torch.arange(seq_len, device=device, dtype=torch.float)
+ freqs = einsum("i , j -> i j", seq, inv_freq)
+ positions = torch.cat((freqs, freqs), dim=-1)
+ pos_sin, pos_cos = positions.sin()[None, None, :, :], positions.cos()[None, None, :, :]
+ self.__cache["rope_pos_sin"] = pos_sin
+ self.__cache["rope_pos_cos"] = pos_cos
+ return pos_sin, pos_cos
+
+ def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
+ B, nh, T, hs = x.size()
+ x = x.view(B, nh, T, 2, hs // 2)
+ x1, x2 = x.unbind(dim=-2)
+ return torch.cat((-x2, x1), dim=-1)
+
+ def apply_rotary_pos_emb(self, pos_sin: torch.Tensor, pos_cos: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+ return ((t * pos_cos) + (self.rotate_half(t) * pos_sin)).to(t.dtype)
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor, q_mask=None) -> Tuple[torch.Tensor, torch.Tensor]:
+ if self.config.rope_full_precision:
+ q_, k_ = q.float(), k.float()
+ else:
+ q_, k_ = q, k
+
+ with torch.autocast(q.device.type, enabled=False):
+ query_len, key_len = q_.shape[-2], k_.shape[-2] # could be different if layer_past not None
+ pos_sin, pos_cos = self.get_rotary_embedding(key_len, q_.device)
+ pos_sin = pos_sin.type_as(q_)
+ pos_cos = pos_cos.type_as(q_)
+ if q_mask is None:
+ q_ = self.apply_rotary_pos_emb(
+ pos_sin[:, :, key_len - query_len : key_len, :],
+ pos_cos[:, :, key_len - query_len : key_len, :],
+ q_,
+ )
+ else:
+ q_ = self.apply_rotary_pos_emb(
+ pos_sin[:, :, q_mask, :],
+ pos_cos[:, :, q_mask, :],
+ q_,
+ )
+ k_ = self.apply_rotary_pos_emb(pos_sin, pos_cos, k_)
+ return q_.type_as(q), k_.type_as(k)
+
+
+class Activation(nn.Module):
+ def __init__(self, config: ModelConfig):
+ super().__init__()
+ self.config = config
+
+ @abstractmethod
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ raise NotImplementedError
+
+ @property
+ @abstractmethod
+ def output_multiplier(self) -> float:
+ raise NotImplementedError
+
+ @classmethod
+ def build(cls, config: ModelConfig) -> "Activation":
+ if config.activation_type == ActivationType.gelu:
+ return cast("Activation", GELU(approximate="none"))
+ elif config.activation_type == ActivationType.relu:
+ return cast("Activation", ReLU(inplace=False))
+ elif config.activation_type == ActivationType.silu:
+ return cast("Activation", SiLU(inplace=False))
+ elif config.activation_type == ActivationType.swiglu:
+ return SwiGLU(config)
+ else:
+ raise NotImplementedError(f"Unknown activation: '{config.activation_type}'")
+
+
+class GELU(nn.GELU):
+ @property
+ def output_multiplier(self) -> float:
+ return 1.0
+
+
+class ReLU(nn.ReLU):
+ @property
+ def output_multiplier(self) -> float:
+ return 1.0
+
+
+class SiLU(nn.SiLU):
+ @property
+ def output_multiplier(self) -> float:
+ return 1.0
+
+
+class SwiGLU(Activation):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, gate = x.chunk(2, dim=-1)
+ return F.silu(gate) * x
+
+ @property
+ def output_multiplier(self) -> float:
+ return 0.5
+
+
+def causal_attention_bias(seq_len: int, device: torch.device) -> torch.FloatTensor:
+ att_bias = torch.triu(
+ torch.ones(seq_len, seq_len, device=device, dtype=torch.float),
+ diagonal=1,
+ )
+ att_bias.masked_fill_(att_bias == 1, torch.finfo(att_bias.dtype).min)
+ return att_bias.view(1, 1, seq_len, seq_len) # type: ignore
+
+
+def get_causal_attention_bias(cache: BufferCache, seq_len: int, device: torch.device) -> torch.Tensor:
+ if (causal_bias := cache.get("causal_attention_bias")) is not None and causal_bias.shape[-1] >= seq_len:
+ if causal_bias.device != device:
+ causal_bias = causal_bias.to(device)
+ cache["causal_attention_bias"] = causal_bias
+ return causal_bias
+ with torch.autocast(device.type, enabled=False):
+ causal_bias = causal_attention_bias(seq_len, device)
+ cache["causal_attention_bias"] = causal_bias
+ return causal_bias
+
+
+def alibi_attention_bias(seq_len: int, config: ModelConfig, device: torch.device) -> torch.FloatTensor:
+ alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.float, device=device).view(1, 1, 1, seq_len)
+
+ # shape: (1, 1, seq_len, seq_len)
+ alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.float, device=device).view(1, 1, seq_len, 1)
+ alibi_bias.abs_().mul_(-1)
+
+ # shape: (n_heads,)
+ m = torch.arange(1, config.n_heads + 1, dtype=torch.float, device=device)
+ m.mul_(config.alibi_bias_max / config.n_heads)
+
+ # shape: (1, n_heads, seq_len, seq_len)
+ return alibi_bias * (1.0 / (2 ** m.view(1, config.n_heads, 1, 1))) # type: ignore
+
+
+class LLaDABlock(nn.Module):
+ def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
+ super().__init__()
+ self.layer_id = layer_id
+ self.config = config
+ self.hidden_size = (
+ config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model
+ )
+ self.__cache = cache
+ assert config.d_model % config.n_heads == 0
+
+ self._activation_checkpoint_fn = None
+
+ # Dropout.
+ self.dropout = Dropout(config.residual_dropout)
+
+ # Layer norms.
+ self.k_norm: Optional[LayerNormBase] = None
+ self.q_norm: Optional[LayerNormBase] = None
+ if config.attention_layer_norm:
+ self.k_norm = LayerNormBase.build(
+ config,
+ size=(config.d_model // config.n_heads) * config.effective_n_kv_heads,
+ elementwise_affine=config.attention_layer_norm_with_affine,
+ )
+ self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine)
+
+ # Activation function.
+ self.act = Activation.build(config)
+ assert (self.act.output_multiplier * self.hidden_size) % 1 == 0
+
+ # Attention output projection.
+ self.attn_out = nn.Linear(
+ config.d_model, config.d_model, bias=config.include_bias, device=config.init_device
+ )
+
+ # Feed-forward output projection.
+ self.ff_out = nn.Linear(
+ int(self.act.output_multiplier * self.hidden_size),
+ config.d_model,
+ bias=config.include_bias,
+ device=config.init_device,
+ )
+ self.ff_out._is_residual = True # type: ignore
+
+ # Rotary embeddings.
+ if self.config.rope:
+ self.rotary_emb = RotaryEmbedding(config, self.__cache)
+
+ self.flash_attn_func = None
+ if config.flash_attention:
+ try:
+ from flash_attn import flash_attn_func # type: ignore
+
+ self.flash_attn_func = flash_attn_func
+ except ModuleNotFoundError:
+ pass
+
+ self.use_cache = False
+ self.init_cache()
+
+ def init_cache(self):
+ self.cache = {"k": {}, "v": {}, "out": {}}
+
+ def caching(self, enable: bool = True):
+ self.use_cache = enable
+ self.init_cache()
+
+ def reset_parameters(self):
+ if self.k_norm is not None:
+ self.k_norm.reset_parameters()
+ if self.q_norm is not None:
+ self.q_norm.reset_parameters()
+ init_weights(
+ self.config,
+ self.attn_out,
+ d=self.config.d_model,
+ layer_id=self.layer_id,
+ type_of_module=ModuleType.out_module,
+ )
+ init_weights(
+ self.config,
+ self.ff_out,
+ d=self.ff_out.in_features,
+ layer_id=self.layer_id,
+ type_of_module=ModuleType.out_module,
+ )
+
+ def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]):
+ if strategy == ActivationCheckpointingStrategy.fine_grained:
+ self._activation_checkpoint_fn = activation_checkpoint_function(self.config)
+ else:
+ self._activation_checkpoint_fn = None
+
+ @classmethod
+ def _cast_attn_bias(cls, bias: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor:
+ target_dtype = input_dtype
+ if bias.device.type == "cuda" and torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ elif bias.device.type == "cpu" and torch.is_autocast_cpu_enabled():
+ target_dtype = torch.get_autocast_cpu_dtype()
+ if bias.dtype != target_dtype:
+ bias = bias.to(target_dtype)
+ ensure_finite_(bias, check_neg_inf=True, check_pos_inf=False)
+ return bias
+
+ def _scaled_dot_product_attention(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ attn_mask: Optional[torch.Tensor] = None,
+ dropout_p: float = 0.0,
+ is_causal: bool = False,
+ ) -> torch.Tensor:
+ if self.flash_attn_func is not None and attn_mask is None:
+ r = self.flash_attn_func(
+ q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=False
+ )
+ return r.transpose(1, 2)
+ else:
+ assert k.size(1) == v.size(1)
+ num_kv_heads = k.size(1)
+ num_q_heads = q.size(1)
+ if num_q_heads != num_kv_heads:
+ assert num_q_heads % num_kv_heads == 0
+ k = k.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads)
+ v = v.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads)
+
+ return F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=None,
+ dropout_p=dropout_p,
+ is_causal=False,
+ )
+
+ def attention(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ attention_bias: Optional[torch.Tensor] = None,
+ layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ to_compute_mask=None,
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+ B, T, C = q.size() # batch size, sequence length, d_model
+ dtype = k.dtype
+
+ if self.q_norm is not None and self.k_norm is not None:
+ q = self.q_norm(q).to(dtype=dtype)
+ k = self.k_norm(k).to(dtype=dtype)
+
+ q = q.view(B, -1, self.config.n_heads, C // self.config.n_heads).transpose(1, 2)
+ k = k.view(B, -1, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2)
+ v = v.view(B, -1, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2)
+
+ if layer_past is not None:
+ past_key, past_value = layer_past
+ k = torch.cat((past_key, k), dim=-2)
+ v = torch.cat((past_value, v), dim=-2)
+
+ if self.config.rope:
+ to_compute_index = (
+ to_compute_mask.nonzero(as_tuple=True)[1] if self.use_cache and to_compute_mask is not None else None
+ )
+ q, k = self.rotary_emb(q, k, q_mask=to_compute_index)
+
+ if attention_bias is not None:
+ attention_bias = self._cast_attn_bias(attention_bias, dtype)
+
+ att = self._scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=None,
+ dropout_p=0.0 if not self.training else self.config.attention_dropout,
+ is_causal=False,
+ )
+
+ att = att.transpose(1, 2).contiguous().view(B, T, C)
+
+ return self.attn_out(att), None
+
+ @abstractmethod
+ def forward(
+ self,
+ x: torch.Tensor,
+ attention_bias: Optional[torch.FloatTensor] = None,
+ layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ use_cache: bool = False,
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+ raise NotImplementedError
+
+ @classmethod
+ def build(cls, layer_id: int, config: ModelConfig, cache: BufferCache) -> "LLaDABlock":
+ if config.block_type == BlockType.sequential:
+ return LLaDASequentialBlock(layer_id, config, cache)
+ elif config.block_type == BlockType.llama:
+ return LLaDALlamaBlock(layer_id, config, cache)
+ else:
+ raise NotImplementedError(f"Unknown block type: '{config.block_type}'")
+
+
+class LLaDASequentialBlock(LLaDABlock):
+ def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
+ super().__init__(layer_id, config, cache)
+ self.attn_norm = LayerNorm.build(config)
+ self.ff_norm = LayerNorm.build(config)
+ head_dim = config.d_model // config.n_heads
+ self.fused_dims = (
+ config.d_model,
+ config.effective_n_kv_heads * head_dim,
+ config.effective_n_kv_heads * head_dim,
+ )
+ self.att_proj = nn.Linear(
+ config.d_model, sum(self.fused_dims), bias=config.include_bias | config.include_qkv_bias, device=config.init_device
+ )
+ self.ff_proj = nn.Linear(
+ config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device
+ )
+
+ def reset_parameters(self):
+ super().reset_parameters()
+ self.attn_norm.reset_parameters()
+ self.ff_norm.reset_parameters()
+ init_weights(
+ self.config, self.att_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module
+ )
+ init_weights(
+ self.config, self.ff_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ attention_bias: Optional[torch.Tensor] = None,
+ layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ use_cache: bool = False,
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+ if self._activation_checkpoint_fn is not None:
+ q, k, v = self.att_proj(self._activation_checkpoint_fn(self.attn_norm, x)).split(
+ self.fused_dims, dim=-1
+ )
+ else:
+ q, k, v = self.att_proj(self.attn_norm(x)).split(self.fused_dims, dim=-1)
+
+ if self._activation_checkpoint_fn is not None:
+ att, cache = self._activation_checkpoint_fn( # type: ignore
+ self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache
+ )
+ else:
+ att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache)
+
+ x = x + self.dropout(att)
+
+ og_x = x
+ if self._activation_checkpoint_fn is not None:
+ x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore
+ else:
+ x = self.ff_norm(x)
+ x = self.ff_proj(x)
+ if self._activation_checkpoint_fn is not None:
+ x = self._activation_checkpoint_fn(self.act, x) # type: ignore
+ else:
+ x = self.act(x)
+ x = self.ff_out(x)
+ x = self.dropout(x)
+ x = og_x + x
+
+ return x, cache
+
+
+class LLaDALlamaBlock(LLaDABlock):
+ def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
+ super().__init__(layer_id, config, cache)
+ self.attn_norm = LayerNorm.build(config)
+ self.ff_norm = LayerNorm.build(config)
+ self.__cache = cache
+
+ head_dim = config.d_model // config.n_heads
+ q_proj_out_dim = config.d_model
+ k_proj_out_dim = config.effective_n_kv_heads * head_dim
+ v_proj_out_dim = config.effective_n_kv_heads * head_dim
+ self.q_proj = nn.Linear(
+ config.d_model, q_proj_out_dim, bias=config.include_bias | config.include_qkv_bias, device=config.init_device
+ )
+ self.k_proj = nn.Linear(
+ config.d_model, k_proj_out_dim, bias=config.include_bias | config.include_qkv_bias, device=config.init_device
+ )
+ self.v_proj = nn.Linear(
+ config.d_model, v_proj_out_dim, bias=config.include_bias | config.include_qkv_bias, device=config.init_device
+ )
+
+ self.ff_proj = nn.Linear(
+ config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device
+ )
+ self.up_proj = nn.Linear(
+ config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device
+ )
+
+ def reset_parameters(self):
+ super().reset_parameters()
+ self.attn_norm.reset_parameters()
+ self.ff_norm.reset_parameters()
+ init_weights(self.config, self.q_proj, d=self.config.d_model, layer_id=None)
+ init_weights(self.config, self.k_proj, d=self.config.d_model, layer_id=None)
+ init_weights(self.config, self.v_proj, d=self.config.d_model, layer_id=None)
+ init_weights(self.config, self.ff_proj, d=self.config.d_model, layer_id=None)
+ init_weights(self.config, self.up_proj, d=self.config.d_model, layer_id=None)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ attention_bias: Optional[torch.Tensor] = None,
+ layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ use_cache: bool = False,
+ cat="cond",
+ to_compute_mask=None,
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+ _B, _T, D = x.shape
+
+ x_normed = self.attn_norm(x)
+ q = self.q_proj(x_normed)
+ k = self.k_proj(x_normed)
+ v = self.v_proj(x_normed)
+
+ if use_cache:
+ if cat not in self.cache["k"]:
+ self.cache["k"][cat] = torch.zeros_like(x)
+ self.cache["v"][cat] = torch.zeros_like(x)
+ if to_compute_mask is not None:
+ self.cache["k"][cat][to_compute_mask] = k.view(-1, D)
+ self.cache["v"][cat][to_compute_mask] = v.view(-1, D)
+ k = self.cache["k"][cat]
+ v = self.cache["v"][cat]
+ else:
+ self.cache["k"][cat] = k
+ self.cache["v"][cat] = v
+
+ if self._activation_checkpoint_fn is not None:
+ att, cache = self._activation_checkpoint_fn( # type: ignore
+ self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache
+ )
+ else:
+ att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask)
+
+ x = x + self.dropout(att)
+
+ og_x = x
+ if self._activation_checkpoint_fn is not None:
+ x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore
+ else:
+ x = self.ff_norm(x)
+ x, x_up = self.ff_proj(x), self.up_proj(x)
+ if self._activation_checkpoint_fn is not None:
+ x = self._activation_checkpoint_fn(self.act, x) # type: ignore
+ else:
+ x = self.act(x)
+ x = x * x_up
+ x = self.ff_out(x)
+ x = self.dropout(x)
+ x = og_x + x
+
+ return x, cache
+
+
+class LLaDAOutput(NamedTuple):
+ logits: torch.FloatTensor
+ attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]]
+ hidden_states: Optional[Tuple[torch.Tensor]]
+
+
+class LLaDAGenerateOutput(NamedTuple):
+ token_ids: torch.LongTensor
+ scores: torch.FloatTensor
+
+
+class LLaDABlockGroup(nn.ModuleList):
+ def __init__(self, config: ModelConfig, layer_offset: int, modules: Optional[Iterable[nn.Module]] = None):
+ super().__init__(modules)
+ self.config = config
+ self.layer_offset = layer_offset
+ self.activation_checkpointing_strategy: Optional[ActivationCheckpointingStrategy] = None
+ self._activation_checkpoint_fn = activation_checkpoint_function(self.config)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ attention_bias: Optional[torch.FloatTensor] = None,
+ layers_past: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+ use_cache: bool = False,
+ ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]:
+ attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None
+ for block_idx, block in enumerate(self):
+ layer_past = None if layers_past is None else layers_past[block_idx]
+ block_idx += self.layer_offset
+ if (
+ (self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.whole_layer)
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_two
+ and block_idx % 2 == 0
+ )
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_three
+ and block_idx % 3 == 0
+ )
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_four
+ and block_idx % 4 == 0
+ )
+ ):
+ x, cache = self._activation_checkpoint_fn( # type: ignore
+ block, x, attention_bias=attention_bias, layer_past=layer_past, use_cache=use_cache
+ )
+ else:
+ x, cache = block(x, attention_bias=attention_bias, layer_past=layer_past, use_cache=use_cache)
+ if attn_key_values is not None:
+ assert cache is not None
+ attn_key_values.append(cache)
+ return x, attn_key_values
+
+ def reset_parameters(self):
+ for block in self:
+ block.reset_parameters()
+
+ def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]):
+ self.activation_checkpointing_strategy = strategy
+ for block in self:
+ block.set_activation_checkpointing(strategy)
+
+
+class LLaDAModel(nn.Module):
+ def __init__(self, config: ModelConfig, init_params: bool = True):
+ super().__init__()
+ self.config = config
+ self.__cache = BufferCache()
+
+ if self.config.alibi and self.config.flash_attention:
+ raise Exception("ALiBi is currently not supported with FlashAttention")
+
+ if self.config.alibi and self.config.rope:
+ raise Exception("ALiBi and RoPE are mutually exclusive")
+
+ if self.config.embedding_size is not None and self.config.embedding_size != self.config.vocab_size:
+ if self.config.embedding_size < self.config.vocab_size:
+ raise Exception("embedding size should be at least as big as vocab size")
+ elif self.config.embedding_size % 128 != 0:
+ import warnings
+
+ warnings.warn(
+ "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning, stacklevel=2
+ )
+
+ self.activation_checkpointing_strategy: Optional[ActivationCheckpointingStrategy] = None
+ self._activation_checkpoint_fn: Callable = activation_checkpoint_function(self.config)
+
+ if not (
+ 0 < self.config.block_group_size <= self.config.n_layers
+ and self.config.n_layers % self.config.block_group_size == 0
+ ):
+ raise Exception("n layers must be divisible by block group size")
+
+ torch.backends.cuda.enable_flash_sdp(True)
+ torch.backends.cuda.enable_mem_efficient_sdp(False) # this is super slow so make sure torch won't use it
+
+ self.transformer = nn.ModuleDict(
+ dict(
+ wte=nn.Embedding(
+ config.embedding_size or config.vocab_size, config.d_model, device=config.init_device
+ ),
+ emb_drop=Dropout(config.embedding_dropout),
+ ln_f=LayerNorm.build(config),
+ )
+ )
+
+ blocks = [LLaDABlock.build(i, config, self.__cache) for i in range(config.n_layers)]
+ if self.config.block_group_size > 1:
+ block_groups = [
+ LLaDABlockGroup(config, i, blocks[i : i + config.block_group_size])
+ for i in range(0, config.n_layers, config.block_group_size)
+ ]
+ self.transformer.update({"block_groups": nn.ModuleList(block_groups)})
+ else:
+ self.transformer.update({"blocks": nn.ModuleList(blocks)})
+
+ if not (self.config.alibi or self.config.rope):
+ self.transformer.update(
+ {"wpe": nn.Embedding(config.max_sequence_length, config.d_model, device=config.init_device)}
+ )
+ if not config.weight_tying:
+ self.transformer.update(
+ {
+ "ff_out": nn.Linear(
+ config.d_model,
+ config.embedding_size or config.vocab_size,
+ bias=config.include_bias,
+ device=config.init_device,
+ )
+ }
+ )
+ if init_params and self.config.init_device != "meta":
+ self.reset_parameters()
+ self.__num_fwd_flops: Optional[int] = None
+
+ if self.config.alibi:
+ get_causal_attention_bias(self.__cache, config.max_sequence_length, _non_meta_init_device(config))
+ self.get_alibi_attention_bias(config.max_sequence_length, _non_meta_init_device(config))
+
+ self.logit_cache = {}
+
+ def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]):
+ self.activation_checkpointing_strategy = strategy
+ if self.config.block_group_size != 1:
+ for block_group in self.transformer.block_groups:
+ block_group.set_activation_checkpointing(strategy)
+ else:
+ for block in self.transformer.blocks:
+ block.set_activation_checkpointing(strategy)
+
+ @property
+ def device(self) -> torch.device:
+ device: torch.device = self.transformer.wte.weight.device # type: ignore
+ if device.type == "meta":
+ return _non_meta_init_device(self.config)
+ else:
+ return device
+
+ def reset_parameters(self):
+ logger.info("Initializing model parameters...")
+ init_weights(
+ self.config,
+ self.transformer.wte, # type: ignore
+ std_factor=(0.5 * math.sqrt(self.config.d_model)) if self.config.scale_logits else 1.0,
+ type_of_module=ModuleType.emb,
+ )
+ if hasattr(self.transformer, "wpe"):
+ init_weights(self.config, self.transformer.wpe, type_of_module=ModuleType.emb) # type: ignore
+
+ self.transformer.ln_f.reset_parameters() # type: ignore
+
+ if hasattr(self.transformer, "ff_out"):
+ init_weights(self.config, self.transformer.ff_out, type_of_module=ModuleType.final_out) # type: ignore
+
+ if self.config.block_group_size == 1:
+ for block in self.transformer.blocks:
+ block.reset_parameters()
+ else:
+ for block_group in self.transformer.block_groups:
+ block_group.reset_parameters()
+
+ def get_alibi_attention_bias(self, seq_len: int, device: torch.device) -> torch.Tensor:
+ if (alibi_bias := self.__cache.get("alibi_attention_bias")) is not None and alibi_bias.shape[
+ -1
+ ] >= seq_len:
+ if alibi_bias.device != device:
+ alibi_bias = alibi_bias.to(device)
+ self.__cache["alibi_attention_bias"] = alibi_bias
+ return alibi_bias
+ with torch.autocast(device.type, enabled=False):
+ alibi_bias = alibi_attention_bias(seq_len, self.config, device)
+ self.__cache["alibi_attention_bias"] = alibi_bias
+ return alibi_bias
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor,
+ input_embeddings: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ attention_bias: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Sequence[Tuple[torch.Tensor, torch.Tensor]]] = None,
+ last_logits_only: bool = False,
+ output_hidden_states: Optional[bool] = None,
+ use_cache=False,
+ to_compute_mask=None,
+ cat="",
+ ) -> LLaDAOutput:
+ if use_cache and to_compute_mask is not None:
+ input_ids = input_ids[to_compute_mask].view(input_ids.shape[0], -1)
+
+ assert not self.config.alibi, "Alibi length extrapolation is not supported for MDM."
+ assert self.config.rope, "Rope must be used in Llama-Encoder for MDM."
+
+ output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+
+ if past_key_values:
+ assert len(past_key_values) == self.config.n_layers
+
+ batch_size, seq_len = input_ids.size() if input_embeddings is None else input_embeddings.size()[:2]
+ if past_key_values is None:
+ past_length = 0
+ else:
+ past_length = past_key_values[0][0].size(-2)
+
+ x = self.transformer.wte(input_ids) if input_embeddings is None else input_embeddings # type: ignore
+
+ if self.config.input_emb_norm:
+ x = x * (self.config.d_model**0.5)
+
+ if not (self.config.alibi or self.config.rope):
+ pos = torch.arange(past_length, past_length + seq_len, dtype=torch.long, device=x.device).unsqueeze(0)
+ pos_emb = self.transformer.wpe(pos) # type: ignore
+ x = pos_emb + x
+
+ x = self.transformer.emb_drop(x) # type: ignore
+
+ if attention_mask is not None and 0.0 in attention_mask:
+ attention_mask = attention_mask.to(dtype=torch.float).view(batch_size, -1)[:, None, None, :]
+ attention_mask = (1.0 - attention_mask) * torch.finfo(attention_mask.dtype).min
+ else:
+ attention_mask = None
+
+ if (
+ attention_bias is not None
+ or attention_mask is not None
+ or self.config.alibi
+ or past_key_values is not None
+ ):
+ if attention_bias is None and self.config.alibi:
+ attention_bias = get_causal_attention_bias(
+ self.__cache, past_length + seq_len, x.device
+ ) + self.get_alibi_attention_bias(past_length + seq_len, x.device)
+ elif attention_bias is None:
+ attention_bias = get_causal_attention_bias(self.__cache, past_length + seq_len, x.device)
+ elif attention_bias.dtype in (torch.int8, torch.bool):
+ attention_bias = attention_bias.to(dtype=torch.float)
+ attention_bias.masked_fill_(attention_bias == 0.0, torch.finfo(attention_bias.dtype).min)
+
+ mask_len = seq_len
+ if attention_mask is not None:
+ mask_len = attention_mask.shape[-1]
+ elif past_key_values is not None:
+ mask_len = past_key_values[0][0].shape[-2] + seq_len
+ attention_bias = attention_bias[:, :, :mask_len, :mask_len].to(dtype=torch.float)
+
+ if attention_mask is not None:
+ attention_bias = attention_bias + attention_mask
+ ensure_finite_(attention_bias, check_neg_inf=True, check_pos_inf=False)
+
+ attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None
+
+ all_hidden_states = []
+
+ if self.config.block_group_size == 1:
+ for block_idx, block in enumerate(self.transformer.blocks):
+ if output_hidden_states:
+ all_hidden_states.append(x)
+
+ layer_past = None if past_key_values is None else past_key_values[block_idx]
+ if (
+ (self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.whole_layer)
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_two
+ and block_idx % 2 == 0
+ )
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_three
+ and block_idx % 3 == 0
+ )
+ or (
+ self.activation_checkpointing_strategy == ActivationCheckpointingStrategy.one_in_four
+ and block_idx % 4 == 0
+ )
+ ):
+ x, _ = self._activation_checkpoint_fn(
+ block, x, attention_bias=attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask, use_cache=use_cache, cat=cat
+ )
+ else:
+ x, _ = block(
+ x, attention_bias=attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask, use_cache=use_cache, cat=cat
+ )
+ else:
+ for group_idx, block_group in enumerate(self.transformer.block_groups):
+ if output_hidden_states:
+ all_hidden_states.append(x)
+
+ layers_past = (
+ None
+ if past_key_values is None
+ else past_key_values[
+ group_idx * self.config.block_group_size : (group_idx + 1) * self.config.block_group_size
+ ]
+ )
+ x, _ = block_group(
+ x, attention_bias=attention_bias, layers_past=layers_past, to_compute_mask=to_compute_mask, use_cache=use_cache, cat=cat
+ )
+
+ if last_logits_only:
+ x = x[:, -1, :].unsqueeze(1)
+
+ x = self.transformer.ln_f(x) # type: ignore
+ if output_hidden_states:
+ all_hidden_states.append(x)
+
+ if self.config.weight_tying:
+ logits = F.linear(x, self.transformer.wte.weight, None) # type: ignore
+ else:
+ logits = self.transformer.ff_out(x) # type: ignore
+ if self.config.scale_logits:
+ logits.mul_(1 / math.sqrt(self.config.d_model))
+
+ if use_cache:
+ if cat not in self.logit_cache:
+ self.logit_cache[cat] = torch.zeros_like(logits)
+ if to_compute_mask is not None:
+ self.logit_cache[cat][to_compute_mask] = logits.view(-1, logits.shape[-1])
+ logits = self.logit_cache[cat]
+ else:
+ self.logit_cache[cat] = logits
+
+ return LLaDAOutput(
+ logits=logits, attn_key_values=attn_key_values, hidden_states=tuple(all_hidden_states) if output_hidden_states else None
+ ) # type: ignore[arg-type]
+
+ def caching(self, enable: bool = True):
+ for block in self.transformer.blocks:
+ block.caching(enable)
+ self.logit_cache = {}
+
+ def empty_cache(self):
+ for block in self.transformer.blocks:
+ block.init_cache()
+ self.logit_cache = {}
+
+
+def create_model_config_from_pretrained_config(config: LLaDAConfig):
+ kwargs = {}
+ for field in fields(ModelConfig):
+ kwargs[field.name] = getattr(config, field.name)
+
+ model_config = ModelConfig(**kwargs)
+ return model_config
+
+
+class LLaDAModelLM(PreTrainedModel):
+ config_class = LLaDAConfig
+ base_model_prefix = "model"
+ _no_split_modules = ["LLaDABlock", "LLaDASequentialBlock", "LLaDALlamaBlock"]
+ # Transformers >= 4.57 may query this during infer_auto_device_map.
+ all_tied_weights_keys: Dict[str, str] = {}
+
+ def __init__(self, config: LLaDAConfig, model: Optional[LLaDAModel] = None, init_params: bool = False):
+ super().__init__(config)
+ self.all_tied_weights_keys = {}
+
+ if not model:
+ model_config = create_model_config_from_pretrained_config(config)
+ model_config.init_device = "cpu"
+ self.model = LLaDAModel(model_config, init_params=init_params)
+ else:
+ self.model = model
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ attention_bias: Optional[torch.Tensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ labels: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ cache_position: Optional[Cache] = None, # This is a hack mitigation of an issue in transformers `4.39.x`
+ use_cache=False,
+ to_compute_mask=None,
+ cat="",
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
+ if output_attentions:
+ raise ValueError("output_attentions is not yet supported in LLaDA")
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.model.forward(
+ input_ids=input_ids,
+ input_embeddings=inputs_embeds,
+ attention_mask=attention_mask,
+ attention_bias=attention_bias,
+ past_key_values=past_key_values,
+ output_hidden_states=output_hidden_states,
+ use_cache=use_cache,
+ to_compute_mask=to_compute_mask,
+ cat=cat,
+ )
+
+ logits = outputs.logits
+ hidden_states = outputs.hidden_states
+
+ loss = None
+ if labels is not None:
+ import warnings
+
+ warnings.warn("Note that for LLaDA, you cannot calculate the loss here.", UserWarning, stacklevel=2)
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ logits=logits,
+ past_key_values=outputs.attn_key_values,
+ hidden_states=hidden_states,
+ )
+
+ def can_generate(self) -> bool:
+ return True
+
+ def prepare_inputs_for_generation(
+ self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple]] = None, **kwargs
+ ):
+ if past_key_values:
+ input_ids = input_ids[:, -1:]
+ model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values}
+
+ model_inputs.update(kwargs)
+ model_inputs["use_cache"] = kwargs.pop("use_cache", self.config.use_cache)
+ return model_inputs
+
+ def get_input_embeddings(self) -> torch.nn.Module:
+ return self.model.transformer.wte
+
+ def set_input_embeddings(self, value: torch.nn.Module):
+ self.model.transformer.wte = value
+
+ def get_output_embeddings(self):
+ if self.config.weight_tying:
+ return self.model.transformer.wte
+ else:
+ return self.model.transformer.ff_out
+
+ def set_output_embeddings(self, value: torch.nn.Module):
+ if self.config.weight_tying:
+ self.model.transformer.wte = value
+ else:
+ self.model.transformer.ff_out = value
+
+ def tie_weights(self, *args, **kwargs):
+ if self.config.weight_tying:
+ self.model.transformer.ff_out = self.model.transformer.wte
+
+ def caching(self, enable: bool = True):
+ self.model.caching(enable)
+
+ def empty_cache(self):
+ self.model.empty_cache()
+
+
+def create_attention_mask(original_lengths, max_tokens, device):
+ batch_size = len(original_lengths)
+ attention_mask = torch.zeros(batch_size, max_tokens, dtype=torch.bool, device=device)
+ for i, length in enumerate(original_lengths):
+ attention_mask[i, :length] = 1
+ return attention_mask
+
+
+class LLaDAForMultiModalGeneration(LLaDAModelLM):
+ config_class = LLaDAConfig
+ base_model_prefix = "model"
+
+ def __init__(self, config: LLaDAConfig, *args, **kwargs):
+ logger.info(f"Initializing MMadaModelLM with config: {config}")
+ super().__init__(config, *args, **kwargs)
+
+ def forward(self, input_ids=None, labels=None, infer=False, use_cache=False, to_compute_mask=None, cat="", **kwargs):
+ input_ids = input_ids.tolist()
+ max_tokens = max([len(_) for _ in input_ids])
+ original_lengths = [len(example) for example in input_ids]
+ input_ids = [example + [0] * (max_tokens - len(example)) for example in input_ids]
+ input_ids = torch.tensor(input_ids, dtype=torch.int64, device=self.device)
+ attention_mask = create_attention_mask(original_lengths, max_tokens, self.device)
+
+ output = LLaDAModelLM.forward(
+ self, input_ids=input_ids, attention_mask=attention_mask, use_cache=use_cache, to_compute_mask=to_compute_mask, cat=cat
+ )
+ if infer:
+ return output
+
+ def get_fsdp_wrap_module_list(self) -> List:
+ modules = [*list(self.model.transformer.blocks), self.model.transformer.ff_out]
+ return modules
+
+
+AutoConfig.register("llada", LLaDAConfig)
+
+# --- End of model definition ---
+
+EXAMPLE_DOC_STRING = """
+Examples:
+```py
+>>> import torch
+>>> from PIL import Image
+>>> from diffusers import VQModel, DiffusionPipeline
+>>> from transformers import AutoTokenizer
+>>> from diffusers.utils import load_image
+
+>>> CHECKPOINT = "Alpha-VLLM/Lumina-DiMOO"
+
+>>> # Load VQ-VAE and tokenizer
+>>> vqvae = VQModel.from_pretrained(CHECKPOINT, subfolder="vqvae").to(device=device, dtype=torch_dtype)
+>>> tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT, trust_remote_code=True)
+
+>>> # Initialize the Lumina-DiMOO pipeline
+>>> pipe = DiffusionPipeline.from_pretrained(
+... CHECKPOINT,
+... custom_pipeline="lumina_dimoo",
+... vqvae=vqvae,
+... tokenizer=tokenizer,
+... torch_dtype=torch.bfloat16
+... )
+>>> pipe.to("cuda")
+
+>>> # Load input image
+>>> input_image = Image.open("path/to/your/ref_image.png").convert("RGB")
+
+>>> prompt = (
+... " your prompt. "
+... )
+
+>>> # Run image-to-image generation
+>>> out = pipe(
+... prompt=prompt,
+... image=input_image,
+... edit_type="depth_control",
+... num_inference_steps=64,
+... task="image_to_image",
+... )
+
+>>> out.images[0].save("i2i_test_output.png")
+"""
+
+
+# --- Helper functions ---
+
+
+def cosine_schedule(t):
+ return torch.cos(t * math.pi / 2)
+
+
+def gumbel_noise(t: torch.Tensor, *, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ if generator is None:
+ u = torch.rand_like(t)
+ else:
+ u = torch.rand(t.shape, device=t.device, dtype=t.dtype, generator=generator)
+ return -torch.log(-torch.log(u + 1e-20) + 1e-20)
+
+
+def add_gumbel_noise(logits, temperature):
+ """
+ Gumbel noise addition function
+ According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality
+ Therefore using float64
+ """
+ if temperature == 0:
+ return logits
+ logits = logits.to(torch.float64)
+ noise = torch.rand_like(logits, dtype=torch.float64)
+ gumbel_noise = (- torch.log(noise)) ** temperature
+ return logits.exp() / gumbel_noise
+
+
+def gumbel_max_sample(logits, temperature=1.0, generator=None):
+ if temperature == 0.0:
+ return logits.argmax(dim=-1)
+ gumbel_noise_ = gumbel_noise(logits, generator=generator)
+ return torch.argmax(logits / temperature + gumbel_noise_, dim=-1)
+
+def get_num_transfer_tokens(mask_index, steps):
+ """
+ In the reverse process, the interval [0, 1] is uniformly discretized into steps intervals
+ Since LLaDA employs a linear noise schedule (as defined in Eq.(8)),
+ the expected number of tokens transitioned at each step should be consistent
+
+ This function is designed to precompute the number of tokens that need to be transitioned at each step
+ """
+ mask_num = mask_index.sum(dim=1, keepdim=True)
+
+ base = mask_num // steps
+ remainder = mask_num % steps
+
+ num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
+
+ for i in range(mask_num.size(0)):
+ num_transfer_tokens[i, :remainder[i]] += 1
+
+ return num_transfer_tokens
+
+
+def mask_by_random_topk(keep_n, probs, temperature=1.0, generator=None):
+ B, _S = probs.shape
+ noise = gumbel_noise(probs, generator=generator)
+
+ conf = probs / temperature + noise
+
+ mask = torch.zeros_like(conf, dtype=torch.bool)
+ for i in range(B):
+ k = keep_n[i]
+ if k > 0:
+ top_k_indices = torch.topk(conf[i], k, largest=True).indices
+ mask[i, top_k_indices] = True
+ return mask
+
+
+def calculate_vq_params(height, width, vae_scale_factor=32):
+ token_grid_height = height // vae_scale_factor
+ token_grid_width = width // vae_scale_factor
+ seq_len = token_grid_height * token_grid_width
+ newline_every = token_grid_width
+ return seq_len, newline_every, token_grid_height, token_grid_width
+
+
+def add_break_line(tokens, token_grid_height, token_grid_width, new_number):
+ new_tokens = []
+ for i in range(token_grid_height):
+ start = i * token_grid_width
+ end = (i + 1) * token_grid_width
+ row = tokens[start:end]
+ new_tokens.extend(row)
+ if i < token_grid_height - 1:
+ new_tokens.append(new_number)
+ return new_tokens
+
+
+def generate_crop_size_list(num_patches, patch_size, max_ratio=4.0):
+ assert max_ratio >= 1.0
+ crop_size_list = []
+ wp, hp = num_patches, 1
+ while wp > 0:
+ if max(wp, hp) / min(wp, hp) <= max_ratio:
+ crop_size_list.append((wp * patch_size, hp * patch_size))
+ if (hp + 1) * wp <= num_patches:
+ hp += 1
+ else:
+ wp -= 1
+ return crop_size_list
+
+
+def center_crop(pil_image, crop_size):
+ while pil_image.size[0] >= 2 * crop_size[0] and pil_image.size[1] >= 2 * crop_size[1]:
+ pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+
+ scale = max(crop_size[0] / pil_image.size[0], crop_size[1] / pil_image.size[1])
+ pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+
+ crop_left = random.randint(0, pil_image.size[0] - crop_size[0])
+ crop_upper = random.randint(0, pil_image.size[1] - crop_size[1])
+ crop_right = crop_left + crop_size[0]
+ crop_lower = crop_upper + crop_size[1]
+ return pil_image.crop(box=(crop_left, crop_upper, crop_right, crop_lower))
+
+
+def var_center_crop(pil_image, crop_size_list, random_top_k=1):
+ w, h = pil_image.size
+ rem_percent = [min(cw / w, ch / h) / max(cw / w, ch / h) for cw, ch in crop_size_list]
+ crop_size = random.choice(
+ sorted(((x, y) for x, y in zip(rem_percent, crop_size_list)), reverse=True)[:random_top_k]
+ )[1]
+ return center_crop(pil_image, crop_size)
+
+
+def preprocess_image(image: Image.Image):
+ image = image.convert("RGB")
+ w, h = image.size
+ w, h = (x - x % 32 for x in (w, h))
+ image = image.resize((w, h), resample=Image.LANCZOS)
+ image = np.array(image).astype(np.float32) / 255.0
+ image = image[None].transpose(0, 3, 1, 2)
+ image = torch.from_numpy(image)
+ return 2.0 * image - 1.0
+
+
+def encode_img_with_breaks(image, vqvae, special_tokens, vae_scale_factor: int = 16):
+ """
+ Encode image, add VQ offset, add newlines, and wrap with BOI/EOI tokens.
+ This function mirrors the logic from the original inference script.
+ """
+ orig = image.convert("RGB")
+ image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor, do_normalize=False)
+ pixels = image_processor.preprocess(orig).to(vqvae.device, dtype=vqvae.dtype)
+ latents = vqvae.encode(pixels).latents
+
+ latents_bsz, _, lat_h, lat_w = latents.shape
+
+ quantized = vqvae.quantize(latents)[2][2] + special_tokens["image_token_offset"]
+ quantized_with_offset = quantized.reshape(latents_bsz, lat_h, lat_w).flatten().tolist()
+
+ tokens_with_breaks = add_break_line(
+ quantized_with_offset, lat_h, lat_w, special_tokens["newline_token"]
+ )
+ return [special_tokens["boi"]] + tokens_with_breaks + [special_tokens["eoi"]]
+
+
+def create_prompt_templates():
+ """Create prompt templates for various tasks based on prompt_utils.py"""
+ templates = {
+ "text_understanding": "You are a multimodal model that can process both text and images. Answer the following question based on the provided images. Analyze each image and combine relevant details to answer.",
+ "image_generation": "Generate an image according to the text prompt.",
+ "image_editing": "Generate an image applying the following editing instruction based on the original image.",
+ "dense_prediction": "Perform dense prediction on the given images.",
+ "control_generation": "Generate an image according to the text prompt and the given control image.",
+ "subject_generation": "Generate an image according to the text prompt and the given object image.",
+ "multi_view": "Generate a view-image based on the given image.",
+ "style_transfer": "Transform the current image into the style of the provided image.",
+ }
+ return templates
+
+
+def generate_image_to_image_prompt(prompt_text, edit_type, templates):
+ """
+ Generate prompt for image-to-image generation based on prompt_utils.py
+ """
+ if "dense" in edit_type or "canny_pred" in edit_type:
+ des = {
+ "canny": "canny edge map",
+ "hed": "hed edge map",
+ "normal": "normal map",
+ "sam2mask": "sam2 mask",
+ "depth": "depth map",
+ "openpose": "pose estimation map",
+ }
+ system_prompt = templates["dense_prediction"]
+ prompt_text_used = f"Generate a {des.get(edit_type.split('_')[0], 'dense map')} according to the image."
+
+ elif "control" in edit_type:
+ system_prompt = templates["control_generation"]
+ prompt_text_used = prompt_text
+
+ elif "subject" in edit_type:
+ system_prompt = templates["subject_generation"]
+ prompt_text_used = prompt_text
+
+ elif "edit" in edit_type:
+ system_prompt = templates["image_editing"]
+ prompt_text_used = prompt_text
+
+ elif "ref_transfer" in edit_type or "image_ref_transfer" in edit_type:
+ system_prompt = templates["style_transfer"]
+ prompt_text_used = "Transform the current image into the style of the provided image."
+
+ elif "multi_view" in edit_type:
+ system_prompt = templates["multi_view"]
+ prompt_text_used = f"Generate the {edit_type.split('_')[-1]} view based on the provided front view."
+
+ else:
+ system_prompt = "Generate an image according to the prompt and image."
+ prompt_text_used = prompt_text
+
+ input_prompt = "" + system_prompt + " " + "" + prompt_text_used + " "
+ uncon_prompt = "" + system_prompt + " " + "" + "" + " "
+
+ return input_prompt, uncon_prompt
+
+def generate_text_to_image_prompt(prompt_text: str, templates: Optional[Dict] = None) -> Tuple[str, str]:
+ """
+ Generate prompt for text-to-image generation
+
+ Args:
+ prompt_text: User input text prompt
+ templates: Optional prompt templates dict
+
+ Returns:
+ Tuple of (input_prompt, unconditional_prompt)
+ """
+ if templates is None:
+ templates = create_prompt_templates()
+
+ system_prompt = templates["image_generation"]
+ input_prompt = "" + system_prompt + " " + "" + prompt_text + " "
+ uncon_prompt = "" + system_prompt + " " + "" + "" + " "
+
+ return input_prompt, uncon_prompt
+
+
+def generate_multimodal_understanding_prompt(question: str, templates: Optional[Dict] = None) -> str:
+ """
+ Generate prompt for multimodal understanding (MMU)
+
+ Args:
+ question: User question about the image
+ templates: Optional prompt templates dict
+
+ Returns:
+ Formatted input prompt
+ """
+ if templates is None:
+ templates = create_prompt_templates()
+
+ system_prompt = "You are a multimodal model that can process both text and images. Answer the following question based on the provided images. Analyze each image and combine relevant details to answer."
+ input_prompt = "" + system_prompt + " " + "" + question + " "
+
+ return input_prompt
+
+
+@torch.no_grad()
+def encode_img_with_paint(
+ img: Image.Image,
+ vqvae: VQModel,
+ *,
+ mask_h_ratio: float = 1, # Height ratio
+ mask_w_ratio: float = 0.2, # Width ratio
+ gray_value: int = 127, # Visualization gray value
+ downsample_mode: str = "area",# Pixel mask alignment to latent grid
+ dilate_latent_k: int = 0, # Optional dilation on latent grid (grid count)
+ mask_mode: str = "inpainting", # "inpainting" | "outpainting"
+ special_tokens
+):
+ """
+ Encode image with mask for inpainting/outpainting tasks
+
+ Args:
+ img: Input PIL image
+ vqvae: VQ-VAE model for encoding
+ mask_h_ratio: Height ratio for mask region (default: 1.0)
+ mask_w_ratio: Width ratio for mask region (default: 0.2)
+ gray_value: Gray value for mask visualization (default: 127)
+ downsample_mode: Downsampling mode for mask alignment ("area", "nearest", "bilinear")
+ dilate_latent_k: Dilation kernel size for latent grid (default: 0)
+ mask_mode: Mask mode - "inpainting" (mask inside) or "outpainting" (mask outside)
+
+ Returns:
+ img_token: List[int] - Token sequence with newlines (126084) inserted at row ends;
+ masked positions = 126336, others = index + 126356
+ vis_img: PIL.Image - Gray mask visualization image (consistent with mask_mode)
+
+ Note:
+ * Encoding uses original image strictly; mask only maps to latent grid to determine
+ which tokens are set to MASK_TOKEN_ID.
+ * mask_mode="inpainting": mask inside rectangle; "outpainting": mask outside rectangle (inverse).
+ """
+
+ assert mask_mode in ("inpainting", "outpainting"), "mask_mode must be 'inpainting' or 'outpainting'"
+
+ # --- 1) Calculate center rectangle and generate visualization ---
+ img = img.convert("RGB")
+ W, H = img.size
+ mh = int(round(H * mask_h_ratio))
+ mw = int(round(W * mask_w_ratio))
+ top = (H - mh) // 2
+ left = (W - mw) // 2
+ bottom = top + mh
+ right = left + mw
+
+ if mask_mode == "inpainting":
+ vis_img = img.copy()
+ draw = ImageDraw.Draw(vis_img)
+ draw.rectangle([left, top, right, bottom], fill=(gray_value, gray_value, gray_value))
+ elif mask_mode == "outpainting": # outpainting
+ bg = Image.new("RGB", (W, H), (gray_value, gray_value, gray_value))
+ crop = img.crop((left, top, right, bottom))
+ bg.paste(crop, (left, top))
+ vis_img = bg
+
+ # --- 2) VQ encoding using original image ---
+ vae_scale_factor = 2 ** (len(vqvae.config.block_out_channels) - 1)
+ image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor, do_normalize=False)
+ x = image_processor.preprocess(img).to(vqvae.device) # 1 x 3 x H' x W'
+ latents = vqvae.encode(x).latents # 1 x C x h x w
+ _, _, lat_h, lat_w = latents.shape
+
+ # Quantization indices
+ quant_pack = vqvae.quantize(latents)
+ indices = quant_pack[2][2].view(1, lat_h, lat_w) # 1 x h x w, long
+
+ # --- 3) Pixel mask -> latent grid mask (aligned with encoding input size) ---
+ Hp, Wp = x.shape[-2:]
+ mask_px = torch.zeros((1, 1, Hp, Wp), dtype=torch.float32, device=vqvae.device)
+ # First generate mask where "rectangle inside=1, outside=0"
+ top_p = int(round(top * Hp / H))
+ left_p = int(round(left * Wp / W))
+ bh_p = int(round(mh * Hp / H))
+ bw_p = int(round(mw * Wp / W))
+ mask_px[:, :, top_p:top_p+bh_p, left_p:left_p+bw_p] = 1.0
+
+ # If outpainting, need to invert (outside=1, inside=0 is the masked region)
+ if mask_mode == "outpainting":
+ mask_px = 1.0 - mask_px
+
+ if downsample_mode not in ("nearest", "area", "bilinear"):
+ downsample_mode = "area"
+ mask_lat = F.interpolate(mask_px, size=(lat_h, lat_w), mode=downsample_mode)
+ mask_lat = (mask_lat > 0.5) if downsample_mode == "area" else (mask_lat >= 0.5)
+ mask_lat = mask_lat[0, 0] # h x w (bool)
+
+ # Optional: latent grid dilation (after inversion is applied)
+ if dilate_latent_k > 0:
+ m = mask_lat.float().unsqueeze(0).unsqueeze(0)
+ ker = 2 * dilate_latent_k + 1
+ m = F.max_pool2d(m, kernel_size=ker, stride=1, padding=dilate_latent_k)
+ mask_lat = (m[0, 0] > 0.5)
+
+ # --- 4) Generate tokens: masked positions=MASK_TOKEN_ID, others=indices+VQ_OFFSET ---
+ idx_flat = indices.view(-1)
+ mask_flat = mask_lat.view(-1)
+ tokens = torch.empty_like(idx_flat)
+ tokens[mask_flat] = special_tokens['mask_token']
+ tokens[~mask_flat] = idx_flat[~mask_flat] + special_tokens['image_token_offset']
+ tokens_list = tokens.tolist()
+
+ # --- 5) Insert newlines (no longer wrapped in /, consistent with current return) ---
+
+ img_token = add_break_line(tokens_list, lat_h, lat_w, special_tokens['newline_token'])
+ return img_token, vis_img
+
+
+class LuminaDiMOOPipelineOutput(BaseOutput):
+ """
+ Output class for the Lumina-DiMOO pipeline.
+
+ Args:
+ images (`List[PIL.Image.Image]` or `np.ndarray`, *optional*):
+ List of generated PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ num_channels)`.
+ text (`str`, *optional*):
+ Generated text from the multimodal understanding task.
+ """
+
+ images: Optional[Union[List[Image.Image], np.ndarray]] = None
+ text: Optional[str] = None
+
+
+class LuminaDiMOOPipeline(DiffusionPipeline):
+ """
+ A unified pipeline for Text-to-Image, Image-to-Image, and Multimodal Understanding
+ using the Lumina-DiMOO model.
+
+ This model was contributed by https://huggingface.co/Alpha-VLLM
+
+ Args:
+ llm ([`LLaDAForMultiModalGeneration`]):
+ The core LLM for multimodal generation, e.g., `LLaDAForMultiModalGeneration`.
+ vqvae ([`VQModel`]):
+ Vector Quantized Variational Auto-Encoder (VQ-VAE) model to encode and decode images to and from discrete
+ latent representations.
+ tokenizer ([`AutoTokenizer`):
+ An `AutoTokenizer` to tokenize text prompts.
+ """
+
+ def __init__(
+ self,
+ vqvae: VQModel,
+ tokenizer: AutoTokenizer,
+ checkpoint: Optional[str] = "Alpha-VLLM/Lumina-DiMOO",
+ torch_dtype: Optional[torch.dtype] = torch.bfloat16,
+ device_map: Optional[str] = "auto",
+ low_cpu_mem_usage: bool = True,
+ ):
+ super().__init__()
+ self.register_modules(
+ vqvae=vqvae,
+ tokenizer=tokenizer,
+ )
+
+ self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1)
+ self.special_tokens = {
+ "mask_token": 126336,
+ "newline_token": 126084,
+ "boa": 126354,
+ "eoa": 126355,
+ "boi": 126349,
+ "eoi": 126350,
+ "image_token_offset": 126356,
+ "uncondition":126351
+ }
+ self.prompt_templates = create_prompt_templates()
+
+ # If checkpoint is not provided, reuse the model path from from_pretrained
+ if checkpoint is None:
+ checkpoint = self._name_or_path
+ raise ValueError("A `checkpoint` path must be provided to load the LLM, either directly or via `from_pretrained`.")
+
+ self.llm = LLaDAForMultiModalGeneration.from_pretrained(
+ checkpoint, torch_dtype=torch_dtype, trust_remote_code=True,
+ device_map=device_map,
+ low_cpu_mem_usage=low_cpu_mem_usage,
+ use_safetensors=True,
+ )
+
+ @staticmethod
+ @torch.no_grad()
+ def generate_i2i(
+ model: LLaDAForMultiModalGeneration,
+ prompt: torch.LongTensor,
+ *,
+ seq_len: int = 1024,
+ newline_every: int = 16,
+ timesteps: int = 18,
+ mask_token_id: int = 126336,
+ newline_id: int = 126084,
+ temperature: float = 1.0,
+ cfg_scale: float = 0.0,
+ cfg_img: float = 0.0,
+ uncon_text: torch.LongTensor,
+ uncon_image: torch.LongTensor,
+ code_start: Optional[int] = None,
+ codebook_size: int = 8192,
+ noise_schedule: Callable[[torch.Tensor], torch.Tensor] = cosine_schedule,
+ text_vocab_size: Optional[int] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.LongTensor:
+ """
+ Image-to-image MaskGit generation (supports CFG for text and image)
+
+ Args:
+ model: Model
+ prompt: Prompt tensor
+ seq_len: Sequence length
+ newline_every: Newline interval per row
+ timesteps: Number of timesteps
+ mask_token_id: Mask token id
+ newline_id: Newline token id
+ temperature: Temperature
+ cfg_scale: Text CFG scale
+ cfg_img: Image CFG scale
+ code_start: Prediction image token satrt index
+ uncon_text: Unconditional text input
+ uncon_image: Unconditional image input
+ codebook_size: Codebook size
+ noise_schedule: Noise schedule function
+ text_vocab_size: Text vocabulary size
+ generator: Random number generator
+
+ Returns:
+ Final VQ codes (1, seq_len)
+ """
+ device = next(model.parameters()).device
+ prompt = prompt.to(device)
+ B, _P = prompt.shape
+ assert B == 1, "batch>1 not supported - wrap in loop if needed"
+
+ x = prompt
+
+ vq_mask = x == mask_token_id
+ unknown_cnt = vq_mask.sum(dim=1, keepdim=True)
+ vq_len = unknown_cnt
+
+ # Infer text vocabulary size
+ if text_vocab_size is None:
+ vocab_total = model(torch.zeros(1, 1, dtype=torch.long, device=device), infer=True).logits.size(-1)
+ text_vocab_size = vocab_total - codebook_size
+ vocab_offset = text_vocab_size
+
+ # for step in range(timesteps):
+ for step in tqdm(range(timesteps), desc="Processing", unit="it"):
+ if unknown_cnt.item() == 0:
+ break
+
+ # Calculate number of tokens to keep (continue masking) this round
+ if step < timesteps - 1:
+ frac = noise_schedule(torch.tensor([(step + 1) / timesteps], device=device))
+ keep_n = (vq_len.float() * frac).floor().clamp_min(1).long()
+ else:
+ keep_n = torch.zeros_like(unknown_cnt)
+
+ # Forward pass (with/without CFG)
+ if cfg_scale > 0 or cfg_img > 0:
+ # CFG text
+ uncond_text = torch.cat((uncon_text.to(x.device), x[:, code_start-2:]), dim=1)
+ uncond_text_vq_mask = torch.cat((torch.zeros((1, uncon_text.size(1)), dtype=torch.bool, device=x.device), vq_mask[:, code_start-2:]), dim=1)
+ # CFG image
+ uncond_img = torch.cat((uncon_image.to(x.device), x[:, code_start-2:]), dim=1)
+ uncond_img_vq_mask = torch.cat((torch.zeros((1, uncon_image.size(1)), dtype=torch.bool, device=x.device), vq_mask[:, code_start-2:]), dim=1)
+
+ cond_logits = model(x, infer=True).logits[:, vq_mask[0], vocab_offset : vocab_offset + codebook_size]
+ uncond_logits_text = model(uncond_text, infer=True).logits[:, uncond_text_vq_mask[0], vocab_offset : vocab_offset + codebook_size]
+ uncond_logits_img = model(uncond_img, infer=True).logits[:, uncond_img_vq_mask[0], vocab_offset : vocab_offset + codebook_size]
+ logits = cond_logits + cfg_scale * (cond_logits - uncond_logits_text) + cfg_img * (cond_logits - uncond_logits_img)
+ else:
+ logits = model(x, infer=True).logits[:, vq_mask[0], vocab_offset : vocab_offset + codebook_size]
+
+ sampled = gumbel_max_sample(logits, temperature, generator=generator)
+ sampled_full = sampled + vocab_offset
+ probs = torch.softmax(logits, dim=-1)
+ conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
+
+ flat_idx = vq_mask.nonzero(as_tuple=False)[:, 1]
+ x.view(-1)[flat_idx] = sampled_full.view(-1)
+
+ conf_map = torch.full_like(x, -math.inf, dtype=probs.dtype)
+ conf_map.view(-1)[flat_idx] = conf.view(-1)
+
+ mask_sel = mask_by_random_topk(keep_n.squeeze(1), conf, temperature=temperature, generator=generator)
+ x.view(-1)[flat_idx[mask_sel.view(-1)]] = mask_token_id
+ vq_mask = x == mask_token_id
+ unknown_cnt = vq_mask.sum(dim=1, keepdim=True)
+
+ # Remove newline tokens
+ vq_ids = x[0, code_start:-2]
+ vq_ids = vq_ids[vq_ids != newline_id].view(1, seq_len)
+ return vq_ids
+
+
+ @staticmethod
+ @torch.no_grad()
+ def generate_image(
+ model: LLaDAForMultiModalGeneration,
+ prompt: torch.LongTensor,
+ *,
+ seq_len: int = 1024,
+ newline_every: int = 16,
+ timesteps: int = 18,
+ mask_token_id: int = 126336,
+ newline_id: int = 126084,
+ temperature: float = 1.0,
+ cfg_scale: float = 0.0,
+ uncon_ids: torch.LongTensor,
+ code_start: Optional[int] = None,
+ codebook_size: int = 8192,
+ noise_schedule: Callable[[torch.Tensor], torch.Tensor] = cosine_schedule,
+ text_vocab_size: Optional[int] = None,
+ generator: Optional[torch.Generator] = None,
+ use_cache=True,
+ cache_ratio=0.9,
+ refresh_interval=5,
+ warmup_ratio=0.3
+ ) -> torch.LongTensor:
+ """
+ MaskGit parallel decoding to generate VQ tokens
+
+ Args:
+ model: Model
+ prompt: Prompt tensor
+ seq_len: Sequence length
+ newline_every: Newline interval per row
+ timesteps: Number of timesteps
+ mask_token_id: Mask token id
+ newline_id: Newline token id
+ temperature: Temperature
+ cfg_scale: CFG scale
+ uncon_ids: Unconditional input
+ code_start: Image token satrt index
+ codebook_size: Codebook size
+ noise_schedule: Noise schedule function
+ text_vocab_size: Text vocabulary size
+ generator: Random number generator
+
+ Returns:
+ Final VQ codes (1, seq_len)
+ """
+
+
+ device = next(model.parameters()).device
+ prompt = prompt.to(device)
+ B, _P = prompt.shape
+ assert B == 1, "batch>1 not supported - wrap in loop if needed"
+
+ x = prompt
+
+ vq_mask = x == mask_token_id
+ unknown_cnt = vq_mask.sum(dim=1, keepdim=True)
+ vq_len = unknown_cnt
+
+ if isinstance(model, LLaDAForMultiModalGeneration):
+ model.caching(use_cache)
+ else: # DDP
+ model.module.caching(use_cache)
+
+ warmup_step = int(timesteps * warmup_ratio)
+ refresh_steps = torch.zeros(timesteps, dtype=torch.bool)
+ for step in range(timesteps):
+ if not use_cache or step <= warmup_step or (step-warmup_step) % refresh_interval == 0:
+ refresh_steps[step] = True
+ compute_ratio = 1 - cache_ratio
+
+ # Infer text vocabulary size
+ if text_vocab_size is None:
+ vocab_total = model(torch.zeros(1, 1, dtype=torch.long, device=device), infer=True).logits.size(-1)
+ text_vocab_size = vocab_total - codebook_size
+ vocab_offset = text_vocab_size
+ cond_to_compute_mask = None
+ uncond_to_compute_mask = None
+
+ # for step in range(timesteps):
+ for step in tqdm(range(timesteps), desc="Processing", unit="it"):
+ if unknown_cnt.item() == 0:
+ break
+
+ # Calculate number of tokens to keep (continue masking) this round
+ if step < timesteps - 1:
+ frac = noise_schedule(torch.tensor([(step + 1) / timesteps], device=device))
+ keep_n = (vq_len.float() * frac).floor().clamp_min(1).long()
+ else:
+ keep_n = torch.zeros_like(unknown_cnt)
+
+ if use_cache and step and refresh_steps[step]:
+ if isinstance(model, LLaDAForMultiModalGeneration):
+ model.empty_cache()
+ else: # DDP
+ model.module.empty_cache()
+
+ # Forward pass (with/without CFG)
+ if cfg_scale > 0:
+ uncond = torch.cat((uncon_ids.to(x.device), x[:, code_start-2:]), axis=1)
+ uncond_vq_mask = torch.cat((torch.zeros((1, uncon_ids.size()[1]), dtype=torch.bool).to(x.device), vq_mask[:, code_start-2:]), axis=1)
+ cond_logits = model(x, infer=True,
+ cat='cond', use_cache=use_cache,
+ to_compute_mask = cond_to_compute_mask if not refresh_steps[step] else None,
+ ).logits[..., vocab_offset : vocab_offset + codebook_size]
+ cond_mask_logits = cond_logits[vq_mask].view(B, -1, codebook_size)
+ uncond_logits = model(uncond, infer=True,
+ cat='uncond', use_cache=use_cache,
+ to_compute_mask = uncond_to_compute_mask if not refresh_steps[step] else None
+ ).logits[..., vocab_offset : vocab_offset + codebook_size]
+ uncond_mask_logits = uncond_logits[uncond_vq_mask].view(B, -1, codebook_size)
+ logits = (1 + cfg_scale) * cond_mask_logits - cfg_scale * uncond_mask_logits
+ else:
+ logits = model(x, infer=True).logits[:, vq_mask[0], vocab_offset : vocab_offset + codebook_size]
+
+ sampled = gumbel_max_sample(logits, temperature, generator=generator)
+ sampled_full = sampled + vocab_offset
+ probs = torch.softmax(logits, dim=-1)
+ conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
+
+ flat_idx = vq_mask.nonzero(as_tuple=False)[:, 1]
+ x.view(-1)[flat_idx] = sampled_full.view(-1)
+
+ conf_map = torch.full_like(x, -math.inf, dtype=probs.dtype)
+ conf_map.view(-1)[flat_idx] = conf.view(-1)
+
+ mask_sel = mask_by_random_topk(keep_n.squeeze(1), conf, temperature=temperature, generator=generator)
+ x.view(-1)[flat_idx[mask_sel.view(-1)]] = mask_token_id
+ vq_mask = x == mask_token_id
+ unknown_cnt = vq_mask.sum(dim=1, keepdim=True)
+
+ if use_cache and step < timesteps - 1 and not refresh_steps[step+1]:
+ cond_conf = cond_logits.max(dim=-1)[0]
+ cond_conf_threshold = torch.quantile(cond_conf.to(torch.float), compute_ratio, dim=-1, keepdim=True)
+ cond_to_compute_mask = cond_conf <= cond_conf_threshold
+
+ uncond_conf = uncond_logits.max(dim=-1)[0]
+ uncond_conf_threshold = torch.quantile(uncond_conf.to(torch.float), compute_ratio, dim=-1, keepdim=True)
+ uncond_to_compute_mask = uncond_conf <= uncond_conf_threshold
+
+ # Remove newline tokens
+ vq_ids = x[0, code_start:-2]
+ vq_ids = vq_ids[vq_ids != newline_id].view(1, seq_len)
+ return vq_ids
+
+
+ @staticmethod
+ @torch.no_grad()
+ def generate_text_understanding(
+ model: LLaDAForMultiModalGeneration,
+ prompt,
+ steps=128,
+ gen_length=128,
+ block_length=128,
+ temperature=0.,
+ cfg_scale=0.,
+ remasking='low_confidence',
+ mask_id=126336,
+ code_start: Optional[int] = None,
+ ):
+ """
+ Text understanding generation function
+
+ Args:
+ model: Mask predictor
+ prompt: Input prompt tensor (1, L)
+ steps: Sampling steps, less than or equal to gen_length
+ gen_length: Generated answer length
+ block_length: Block length, less than or equal to gen_length
+ temperature: Categorical distribution sampling temperature
+ cfg_scale: Unsupervised classifier-free guidance scale
+ remasking: Remasking strategy 'low_confidence' or 'random'
+ mask_id: The token id of [MASK] is 126336
+ code_start: Prediction text token satrt index
+ """
+ x = prompt
+
+ prompt_index = x != mask_id
+
+ assert gen_length % block_length == 0
+ num_blocks = gen_length // block_length
+
+ assert steps % num_blocks == 0
+ steps = steps // num_blocks
+
+ for num_block in range(num_blocks):
+ block_mask_index = (x[:, code_start + num_block * block_length: code_start + (num_block + 1) * block_length:] == mask_id)
+ num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
+
+ for i in range(steps):
+ mask_index = x == mask_id
+ if cfg_scale > 0.:
+ un_x = x.clone()
+ un_x[prompt_index] = mask_id
+ x_ = torch.cat([x, un_x], dim=0)
+ logits = model(x_, infer=True).logits
+ logits, un_logits = torch.chunk(logits, 2, dim=0)
+ logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
+ else:
+ logits = model(x, infer=True).logits
+
+ logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+ x0 = torch.argmax(logits_with_noise, dim=-1) # b, l
+
+ if remasking == 'low_confidence':
+ p = F.softmax(logits.to(torch.float64), dim=-1)
+ x0_p = torch.squeeze(
+ torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
+ elif remasking == 'random':
+ x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+ else:
+ raise NotImplementedError(remasking)
+
+ x0_p[:, code_start + (num_block + 1) * block_length:] = -np.inf
+
+ x0 = torch.where(mask_index, x0, x)
+ confidence = torch.where(mask_index, x0_p, -np.inf)
+
+ transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+ for j in range(confidence.shape[0]):
+ _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j, i])
+ transfer_index[j, select_index] = True
+ x[transfer_index] = x0[transfer_index]
+
+
+ return x
+
+
+
+ @torch.no_grad()
+ def _image_to_image(
+ self,
+ prompt: str,
+ image: Union[Image.Image, str],
+ ref_image: Optional[PipelineImageInput] = None,
+ edit_type: str = "canny_pred",
+ num_inference_steps: int = 64,
+ temperature: float = 1.0,
+ cfg_scale: float = 2.5,
+ cfg_img: float = 4.0,
+ output_type: Optional[str] = "pil",
+ ):
+
+ if isinstance(prompt, list):
+ raise ValueError("Batching is not supported for this pipeline.")
+
+ if isinstance(prompt, list):
+ prompt = prompt[0]
+ if isinstance(image, str):
+ image = Image.open(image).convert("RGB")
+ if isinstance(ref_image, str):
+ ref_image = Image.open(ref_image).convert("RGB")
+
+ input_prompt, uncon_text = generate_image_to_image_prompt(prompt, edit_type, self.prompt_templates)
+
+ crop_size_list = generate_crop_size_list((512 // 32) ** 2, 32)
+
+ # Correctly encode input images with newline tokens
+ if "image_ref_transfer" in edit_type:
+ if ref_image is None:
+ raise ValueError("`ref_image` must be provided for `image_ref_transfer` edit type.")
+ processed_img = var_center_crop(image, crop_size_list=crop_size_list)
+ input_img_token = encode_img_with_breaks(processed_img, self.vqvae, self.special_tokens)
+
+ referring_img = var_center_crop(ref_image, crop_size_list=crop_size_list)
+ referring_img_token = encode_img_with_breaks(referring_img, self.vqvae, self.special_tokens)
+
+ image_width, image_height = referring_img.size
+ seq_len, newline_every, token_grid_height, token_grid_width = calculate_vq_params(
+ referring_img.height, referring_img.width, self.vae_scale_factor
+ )
+ else:
+ processed_img = var_center_crop(image, crop_size_list=crop_size_list)
+ input_img_token = encode_img_with_breaks(processed_img, self.vqvae, self.special_tokens)
+ image_width, image_height = processed_img.size
+ seq_len, newline_every, token_grid_height, token_grid_width = calculate_vq_params(
+ processed_img.height, processed_img.width, self.vae_scale_factor
+ )
+
+ prompt_ids = self.tokenizer(input_prompt)["input_ids"]
+ uncon_text_ids = self.tokenizer(uncon_text)["input_ids"]
+
+ img_mask_token = add_break_line(
+ [self.special_tokens["mask_token"]] * seq_len,
+ token_grid_height,
+ token_grid_width,
+ new_number=self.special_tokens["newline_token"],
+ )
+ img_pred_token = (
+ [self.special_tokens["boa"]]
+ + [self.special_tokens["boi"]]
+ + img_mask_token
+ + [self.special_tokens["eoi"]]
+ + [self.special_tokens["eoa"]]
+ )
+
+ if "image_ref_transfer" in edit_type:
+ con_input = prompt_ids[:-1] + input_img_token + referring_img_token + prompt_ids[-1:]
+ uncon_input_text = uncon_text_ids[:-1] + input_img_token + referring_img_token + uncon_text_ids[-1:]
+ else:
+ con_input = prompt_ids[:-1] + input_img_token + prompt_ids[-1:]
+ uncon_input_text = uncon_text_ids[:-1] + input_img_token + uncon_text_ids[-1:]
+ uncon_input_image = prompt_ids
+
+ code_start = len(con_input) + 2
+
+ con_input = torch.tensor(con_input + img_pred_token, device=self.device).unsqueeze(0)
+ uncon_input_text = torch.tensor(uncon_input_text, device=self.device).unsqueeze(0)
+ uncon_input_image = torch.tensor(uncon_input_image, device=self.device).unsqueeze(0)
+
+ vq_tokens = self.generate_i2i(
+ self.llm,
+ con_input,
+ seq_len=seq_len,
+ newline_every=newline_every,
+ timesteps=num_inference_steps,
+ temperature=temperature,
+ cfg_scale=cfg_scale,
+ cfg_img=cfg_img,
+ uncon_text=uncon_input_text,
+ uncon_image=uncon_input_image,
+ code_start=code_start
+ )
+
+ if vq_tokens.shape[1] != token_grid_height * token_grid_width:
+ raise ValueError(
+ f"VQ codes length mismatch: {vq_tokens.shape[1]} != {token_grid_height * token_grid_width} "
+ f"for image size ({image_height},{image_width}) with scale {self.vae_scale_factor}"
+ )
+
+ latents = (
+ vq_tokens.view(1, token_grid_height, token_grid_width).to(self.vqvae.device) - self.special_tokens["image_token_offset"]
+ ).long()
+
+ shape = (1, token_grid_height, token_grid_width, self.vqvae.config.latent_channels)
+
+ recon = self.vqvae.decode(
+ latents,
+ force_not_quantize=True,
+ shape=shape,
+ ).sample.clip(0, 1)
+
+ img_proc = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False)
+ image = img_proc.postprocess(recon.detach(), output_type=output_type)
+
+ return image
+
+ @torch.no_grad()
+ def _text_to_image(
+ self,
+ prompt: str,
+ height: int,
+ width: int,
+ painting_mode: Optional[str] = None,
+ painting_image: Optional[PipelineImageInput] = None,
+ cfg_scale: float = 4.0,
+ use_cache: bool = True,
+ cache_ratio: float = 0.9,
+ refresh_interval: int = 5,
+ warmup_ratio: float = 0.3,
+ num_inference_steps: int = 64,
+ temperature: float = 1.0,
+ mask_h_ratio: float = 1.0,
+ mask_w_ratio: float = 0.2
+ ):
+ if isinstance(painting_image, str):
+ painting_image = Image.open(painting_image)
+
+ if painting_mode and painting_image:
+ width, height = painting_image.size
+
+ seq_len, newline_every, token_grid_height, token_grid_width = calculate_vq_params(height, width, self.vae_scale_factor)
+
+ input_prompt, uncon_prompt = generate_text_to_image_prompt(prompt, self.prompt_templates)
+
+ con_prompt_token = self.tokenizer(input_prompt)["input_ids"]
+ uncon_prompt_token = self.tokenizer(uncon_prompt)["input_ids"]
+
+ if painting_mode:
+ img_mask_token, _img_vis = encode_img_with_paint(
+ painting_image,
+ vqvae=self.vqvae,
+ mask_h_ratio=mask_h_ratio,
+ mask_w_ratio=mask_w_ratio,
+ mask_mode=painting_mode,
+ special_tokens=self.special_tokens,
+ )
+ else:
+ img_mask_token = add_break_line(
+ [self.special_tokens["mask_token"]] * seq_len,
+ token_grid_height,
+ token_grid_width,
+ new_number=self.special_tokens["newline_token"],
+ )
+
+ img_pred_token = (
+ [self.special_tokens["boa"]]
+ + [self.special_tokens["boi"]]
+ + img_mask_token
+ + [self.special_tokens["eoi"]]
+ + [self.special_tokens["eoa"]]
+ )
+
+ prompt_ids = torch.tensor(con_prompt_token + img_pred_token, device=self.device).unsqueeze(0)
+ uncon_ids = torch.tensor(uncon_prompt_token, device=self.device).unsqueeze(0)
+
+ code_start = len(con_prompt_token) + 2
+
+ vq_tokens = self.generate_image(
+ model=self.llm,
+ prompt=prompt_ids,
+ seq_len=seq_len,
+ newline_every=newline_every,
+ timesteps=num_inference_steps,
+ temperature=temperature,
+ cfg_scale=cfg_scale,
+ uncon_ids=uncon_ids,
+ code_start=code_start,
+ use_cache=use_cache,
+ cache_ratio=cache_ratio,
+ refresh_interval=refresh_interval,
+ warmup_ratio=warmup_ratio
+ )
+
+ latents = (
+ vq_tokens.view(1, token_grid_height, token_grid_width).to(self.vqvae.device) - self.special_tokens["image_token_offset"]
+ ).long()
+
+ shape = (1, token_grid_height, token_grid_width, self.vqvae.config.latent_channels)
+ recon = self.vqvae.decode(latents, force_not_quantize=True, shape=shape).sample.clip(0, 1)
+
+ img_proc = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False)
+ image = img_proc.postprocess(recon.detach(), output_type="pil")
+
+ return image
+
+ @torch.no_grad()
+ def _multimodal_understanding(
+ self,
+ prompt: str,
+ image: PipelineImageInput,
+ num_inference_steps: int = 128,
+ gen_length: int = 1024,
+ block_length: int = 128,
+ temperature: float = 0.0,
+ cfg_scale: float = 0.0,
+ remasking: str = "low_confidence",
+ ):
+
+ if isinstance(image, str):
+ image = Image.open(image)
+
+ input_prompt = generate_multimodal_understanding_prompt(prompt)
+ input_ids = self.tokenizer(input_prompt)["input_ids"]
+
+ crop_size_list = generate_crop_size_list((1024 // 32) ** 2, 32)
+ processed_image = var_center_crop(image, crop_size_list=crop_size_list)
+
+ image_width, image_height = processed_image.size
+ _seq_len, _newline_every, _token_grid_height, _token_grid_width = calculate_vq_params(
+ image_height, image_width, self.vae_scale_factor
+ )
+
+ input_img_token = encode_img_with_breaks(processed_image, self.vqvae, self.special_tokens)
+
+ input_token = input_ids[:-1] + input_img_token + input_ids[-1:]
+ code_start = len(input_token) + 1
+
+ input_token = input_token + [self.special_tokens["boa"]] + gen_length * [self.special_tokens["mask_token"]] + [self.special_tokens["eoa"]]
+ input_ids = torch.tensor(input_token, device=self.device).unsqueeze(0)
+
+ output_tokens = self.generate_text_understanding(
+ model=self.llm,
+ prompt=input_ids,
+ steps=num_inference_steps,
+ gen_length=gen_length,
+ block_length=block_length,
+ cfg_scale=cfg_scale,
+ temperature=temperature,
+ remasking=remasking,
+ code_start=code_start
+ )
+
+ generated_text = self.tokenizer.batch_decode(output_tokens[:, code_start:-1], skip_special_tokens=True)[0]
+ return generated_text
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: str,
+ image: Optional[PipelineImageInput] = None,
+ task: str = "auto",
+ **kwargs,
+ ) -> LuminaDiMOOPipelineOutput:
+ r"""
+ Unified entry for 'text_to_image' | 'image_to_image' | 'multimodal_understanding'.
+
+ Examples:
+ {EXAMPLE_DOC_STRING}
+ """
+ if isinstance(prompt, list):
+ prompt = prompt[0]
+
+ if task == "auto":
+ if image is None:
+ task = "text_to_image"
+ elif "edit_type" in kwargs:
+ task = "image_to_image"
+ else:
+ task = "multimodal_understanding"
+
+ if task == "text_to_image":
+ # Default values from inference_t2i.py
+ t2i_kwargs = {
+ "height": kwargs.pop("height", 1024),
+ "width": kwargs.pop("width", 1024),
+ "num_inference_steps": kwargs.pop("num_inference_steps", 64),
+ "cfg_scale": kwargs.pop("cfg_scale", 4.0),
+ "temperature": kwargs.pop("temperature", 1.0),
+ "painting_mode": kwargs.pop("painting_mode", None),
+ "painting_image": kwargs.pop("painting_image", None),
+ "mask_h_ratio": kwargs.pop("mask_h_ratio", 1.0),
+ "mask_w_ratio": kwargs.pop("mask_w_ratio", 0.2),
+ "use_cache": kwargs.pop("use_cache", True),
+ "cache_ratio": kwargs.pop("cache_ratio", 0.9),
+ "refresh_interval": kwargs.pop("refresh_interval", 5),
+ "warmup_ratio": kwargs.pop("warmup_ratio", 0.3),
+ }
+ images = self._text_to_image(prompt=prompt, **t2i_kwargs)
+ return LuminaDiMOOPipelineOutput(images=images, text=None)
+
+ elif task == "image_to_image":
+ if image is None:
+ raise ValueError("`image` must be provided for image_to_image task.")
+ i2i_kwargs = {
+ "ref_image": kwargs.pop("ref_image", None),
+ "edit_type": kwargs.pop("edit_type", "canny_pred"),
+ "num_inference_steps": kwargs.pop("num_inference_steps", 64),
+ "temperature": kwargs.pop("temperature", 1.0),
+ "cfg_scale": kwargs.pop("cfg_scale", 2.5),
+ "cfg_img": kwargs.pop("cfg_img", 4.0),
+ }
+ images = self._image_to_image(prompt=prompt, image=image, **i2i_kwargs)
+ return LuminaDiMOOPipelineOutput(images=images, text=None)
+
+ elif task == "multimodal_understanding":
+ if image is None:
+ raise ValueError("`image` must be provided for multimodal_understanding task.")
+ mmu_kwargs = {
+ "num_inference_steps": kwargs.pop("num_inference_steps", 128),
+ "gen_length": kwargs.pop("gen_length", 1024),
+ "block_length": kwargs.pop("block_length", 256),
+ "temperature": kwargs.pop("temperature", 0.0),
+ "cfg_scale": kwargs.pop("cfg_scale", 0.0),
+ "remasking": kwargs.pop("remasking", "low_confidence"),
+ }
+ text = self._multimodal_understanding(prompt=prompt, image=image, **mmu_kwargs)
+ return LuminaDiMOOPipelineOutput(images=None, text=text)
+
+ else:
+ raise ValueError(f"Unknown task: {task}. Supported tasks are 'text_to_image', 'image_to_image', 'multimodal_understanding', and 'auto'.")
diff --git a/pipelines/lumina_dimmo/notes.md b/pipelines/lumina_dimmo/notes.md
new file mode 100644
index 000000000..1d6f232a0
--- /dev/null
+++ b/pipelines/lumina_dimmo/notes.md
@@ -0,0 +1,163 @@
+
+To load a custom pipeline you just need to pass the `custom_pipeline` argument to `DiffusionPipeline`, as one of the files in `diffusers/examples/community`. Feel free to send a PR with your own pipelines, we will merge them quickly.
+
+
+## Example Usage
+
+# Lumina-DiMOO
+[Project](https://synbol.github.io/Lumina-DiMOO/) / [GitHub](https://github.com/Alpha-VLLM/Lumina-DiMOO/) / [Model](https://huggingface.co/Alpha-VLLM/Lumina-DiMOO)
+
+Lumina-DiMOO is a discrete-diffusion omni-modal foundation model unifying generation and understanding. This implementation integrates a Lumina-DiMOO switch for T2I, I2I editing, and MMU.
+
+#### Key features
+
+- **Unified Discrete Diffusion Architecture**: Employs a fully discrete diffusion framework to process inputs and outputs across diverse modalities.
+- **Versatile Multimodal Capabilities**: Supports a wide range of multimodal tasks, including text-to-image generation (arbitrary and high-resolution), image-to-image generation (e.g., image editing, subject-driven generation, inpainting), and advanced image understanding.
+- **Higher Sampling Efficiency**: Outperforms previous autoregressive (AR) or hybrid AR-diffusion models with significantly faster sampling. A custom caching mechanism further boosts sampling speed by up to 2×.
+
+
+### Example Usage
+
+The Lumina-DiMOO pipeline provides three core functions — T2I, I2I, and MMU.
+For detailed implementation examples and creative applications, please visit the [GitHub](https://github.com/Alpha-VLLM/Lumina-DiMOO)
+
+
+#### Text-to-Image
+**prompt** | **image**
+:-------------------------:|:-------------------------:
+| "A striking photograph of a glass of orange juice on a wooden kitchen table, capturing a playful moment. The orange juice splashes out of the glass and forms the word \"Smile\" in a whimsical, swirling script just above the glass. The background is softly blurred, revealing a cozy, homely kitchen with warm lighting and a sense of comfort." |
+
+```python
+import torch
+
+from diffusers import VQModel, DiffusionPipeline
+from transformers import AutoTokenizer
+
+vqvae = VQModel.from_pretrained("Alpha-VLLM/Lumina-DiMOO", subfolder="vqvae").to(device='cuda', dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("Alpha-VLLM/Lumina-DiMOO", trust_remote_code=True)
+
+pipe = DiffusionPipeline.from_pretrained(
+ "Alpha-VLLM/Lumina-DiMOO",
+ vqvae=vqvae,
+ tokenizer=tokenizer,
+ torch_dtype=torch.bfloat16,
+ custom_pipeline="lumina_dimoo",
+)
+pipe.to("cuda")
+
+prompt = '''A striking photograph of a glass of orange juice on a wooden kitchen table, capturing a playful moment. The orange juice splashes out of the glass and forms the word \"Smile\" in a whimsical, swirling script just above the glass. The background is softly blurred, revealing a cozy, homely kitchen with warm lighting and a sense of comfort.'''
+
+img = pipe(
+ prompt=prompt,
+ task="text_to_image",
+ height=768,
+ width=1536,
+ num_inference_steps=64,
+ cfg_scale=4.0,
+ use_cache=True,
+ cache_ratio=0.9,
+ warmup_ratio=0.3,
+ refresh_interval=5
+).images[0]
+
+img.save("t2i_test_output.png")
+```
+
+#### Image-to-Image
+**prompt** | **image_before** | **image_after**
+:-------------------------:|:-------------------------:|:-------------------------:
+| "A functional wooden printer stand.Nestled next to a brick wall in a bustling city street, it stands firm as pedestrians hustle by, illuminated by the warm glow of vintage street lamps." | | |
+
+```python
+import torch
+
+from diffusers import VQModel, DiffusionPipeline
+from transformers import AutoTokenizer
+from diffusers.utils import load_image
+
+vqvae = VQModel.from_pretrained("Alpha-VLLM/Lumina-DiMOO", subfolder="vqvae").to(device='cuda', dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("Alpha-VLLM/Lumina-DiMOO", trust_remote_code=True)
+
+pipe = DiffusionPipeline.from_pretrained(
+ "Alpha-VLLM/Lumina-DiMOO",
+ vqvae=vqvae,
+ tokenizer=tokenizer,
+ torch_dtype=torch.bfloat16,
+ custom_pipeline="lumina_dimoo",
+)
+pipe.to("cuda")
+
+input_image = load_image(
+ "https://raw.githubusercontent.com/Alpha-VLLM/Lumina-DiMOO/main/examples/example_2.jpg"
+).convert("RGB")
+
+prompt = "A functional wooden printer stand.Nestled next to a brick wall in a bustling city street, it stands firm as pedestrians hustle by, illuminated by the warm glow of vintage street lamps."
+
+img = pipe(
+ prompt=prompt,
+ image=input_image,
+ edit_type="depth_control",
+ num_inference_steps=64,
+ temperature=1.0,
+ cfg_scale=2.5,
+ cfg_img=4.0,
+ task="image_to_image"
+).images[0]
+
+img.save("i2i_test_output.png")
+
+```
+
+
+#### Multimodal Understanding
+**question** | **image** | **answer**
+:-------------------------:|:-------------------------:|:-------------------------:
+| "Please describe the image." | | "The image shows a vibrant orange sports car parked in a showroom. The car has a sleek, aerodynamic design with a prominent front grille and side vents. The body is adorned with black and orange racing stripes, creating a striking contrast against the orange paint. The car is equipped with black alloy wheels and a low-profile body style. The background features a white wall with a large emblem that reads "BREITZEN" and includes a silhouette of a horse and text. The floor is tiled with dark tiles, and the showroom is well-lit, highlighting the car. The overall setting suggests a high-end, possibly luxury, automotive environment."|
+
+
+```python
+import torch
+result.images[0].save(f"flux_fill_controlnet_inpaint_depth{timestamp}.jpg")
+```
+
+
+```python
+import torch
+
+from diffusers import VQModel, DiffusionPipeline
+from transformers import AutoTokenizer
+from diffusers.utils import load_image
+
+vqvae = VQModel.from_pretrained("Alpha-VLLM/Lumina-DiMOO", subfolder="vqvae").to(device='cuda', dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("Alpha-VLLM/Lumina-DiMOO", trust_remote_code=True)
+
+pipe = DiffusionPipeline.from_pretrained(
+ "Alpha-VLLM/Lumina-DiMOO",
+ vqvae=vqvae,
+ tokenizer=tokenizer,
+ torch_dtype=torch.bfloat16,
+ custom_pipeline="lumina_dimoo",
+)
+pipe.to("cuda")
+
+question = "Please describe the image."
+
+input_image = load_image(
+ "https://raw.githubusercontent.com/Alpha-VLLM/Lumina-DiMOO/main/examples/example_8.png"
+).convert("RGB")
+
+out = pipe(
+ prompt=question,
+ image=input_image,
+ task="multimodal_understanding",
+ num_inference_steps=128,
+ gen_length=128,
+ block_length=32,
+ temperature=0.0,
+ cfg_scale=0.0,
+)
+
+text = getattr(out, "text", out)
+with open("mmu_answer.txt", "w", encoding="utf-8") as f:
+ f.write(text.strip() + "\n")
+```
diff --git a/pipelines/lumina_dimmo/pipelines.py b/pipelines/lumina_dimmo/pipelines.py
new file mode 100644
index 000000000..cb264d3a8
--- /dev/null
+++ b/pipelines/lumina_dimmo/pipelines.py
@@ -0,0 +1,58 @@
+from pipelines.lumina_dimmo.lumina_dimoo import LuminaDiMOOPipeline
+from modules.logger import log
+
+
+class LuminaDiMOOTextPipeline(LuminaDiMOOPipeline):
+ def __call__(self, prompt: list[str], height: int = 1024, width: int = 1024, num_inference_steps: int = 64, cfg_scale: float = 4.0, temperature: float = 1.0, painting_mode=None, painting_image=None, mask_h_ratio: float = 1.0, mask_w_ratio: float = 0.2, use_cache: bool = True, cache_ratio: float = 0.9, refresh_interval: int = 5, warmup_ratio: float = 0.3, **kwargs):
+ if isinstance(prompt, list):
+ prompt = prompt[0].strip()
+ task = "text_to_image"
+ log.debug(f'Base args: task={task} height={height} width={width} steps={num_inference_steps} cfg_scale={cfg_scale} temperature={temperature} painting_mode={painting_mode} mask_h_ratio={mask_h_ratio} mask_w_ratio={mask_w_ratio} use_cache={use_cache} cache_ratio={cache_ratio} refresh_interval={refresh_interval} warmup_ratio={warmup_ratio}')
+ return super().__call__(
+ prompt=prompt,
+ task=task,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ cfg_scale=cfg_scale,
+ temperature=temperature,
+ painting_mode=painting_mode,
+ painting_image=painting_image,
+ mask_h_ratio=mask_h_ratio,
+ mask_w_ratio=mask_w_ratio,
+ use_cache=use_cache,
+ cache_ratio=cache_ratio,
+ refresh_interval=refresh_interval,
+ warmup_ratio=warmup_ratio,
+ **kwargs)
+
+
+class LuminaDiMOOImagePipeline(LuminaDiMOOPipeline):
+ def __call__(self, prompt: list[str], image, num_inference_steps: int = 64, temperature: float = 1.0, cfg_scale: float = 2.5, cfg_img: float = 4.0, **kwargs):
+ edit_types = ['dense', 'canny_pred', 'control', 'subject', 'edit', 'ref_transfer', 'multi_view']
+ if isinstance(prompt, list):
+ prompt = prompt[0].strip()
+ task = "image_to_image"
+ edit_type = 'default'
+ ref_image = None
+ for et in edit_types:
+ if prompt.startswith(et):
+ edit_type = et
+ break
+ if isinstance(image, list):
+ if len(image) > 1:
+ ref_image = image[1]
+ if len(image) > 0:
+ image = image[0]
+ log.debug(f'Base args: task={task} edit_type={edit_type} steps={num_inference_steps} cfg_scale={cfg_scale} cfg_img={cfg_img} temperature={temperature} image={image} ref_image={ref_image}')
+ return super().__call__(
+ prompt=prompt,
+ task=task,
+ image=image,
+ ref_image=ref_image,
+ edit_type=edit_type,
+ num_inference_steps=num_inference_steps,
+ temperature=temperature,
+ cfg_scale=cfg_scale,
+ cfg_img=cfg_img,
+ **kwargs)
diff --git a/pipelines/meissonic/pipeline_img2img.py b/pipelines/meissonic/pipeline_img2img.py
index 2aaf9d987..8ff9c0705 100644
--- a/pipelines/meissonic/pipeline_img2img.py
+++ b/pipelines/meissonic/pipeline_img2img.py
@@ -276,7 +276,7 @@ class MeissonicImg2ImgPipeline(DiffusionPipeline):
self.vqvae.float()
latents = self.vqvae.encode(image.to(dtype=self.vqvae.dtype, device=self._execution_device)).latents
- latents_bsz, channels, latents_height, latents_width = latents.shape
+ latents_bsz, _channels, latents_height, latents_width = latents.shape
latents = self.vqvae.quantize(latents)[2][2].reshape(latents_bsz, latents_height, latents_width)
latents = self.scheduler.add_noise(
latents, self.scheduler.timesteps[start_timestep_idx - 1], generator=generator
diff --git a/pipelines/meissonic/pipeline_inpaint.py b/pipelines/meissonic/pipeline_inpaint.py
index aa352d9b4..f43290d3a 100644
--- a/pipelines/meissonic/pipeline_inpaint.py
+++ b/pipelines/meissonic/pipeline_inpaint.py
@@ -289,7 +289,7 @@ class MeissonicInpaintPipeline(DiffusionPipeline):
self.vqvae.float()
latents = self.vqvae.encode(image.to(dtype=self.vqvae.dtype, device=self._execution_device)).latents
- latents_bsz, channels, latents_height, latents_width = latents.shape
+ latents_bsz, _channels, latents_height, latents_width = latents.shape
latents = self.vqvae.quantize(latents)[2][2].reshape(latents_bsz, latents_height, latents_width)
mask = self.mask_processor.preprocess(
diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py
index 2f3149d25..fa7b3b6c2 100644
--- a/pipelines/model_anima.py
+++ b/pipelines/model_anima.py
@@ -81,6 +81,12 @@ def load_anima(checkpoint_info, diffusers_load_config=None):
AnimaTextToImagePipeline = pipeline_mod.AnimaTextToImagePipeline
AnimaLLMAdapter = adapter_mod.AnimaLLMAdapter
+ from pipelines.anima.anima_image import build_anima_pipeline_classes
+ AnimaImageToImagePipeline, AnimaInpaintPipeline = build_anima_pipeline_classes(AnimaTextToImagePipeline)
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["anima"] = AnimaTextToImagePipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["anima"] = AnimaImageToImagePipeline
+ diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["anima"] = AnimaInpaintPipeline
+
# UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a
# community file that bundles both the transformer and the llm_adapter.
transformer, llm_adapter = load_transformer_components(repo_id, diffusers_load_config, AnimaLLMAdapter)
@@ -119,6 +125,9 @@ def load_anima(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ # generic.load_vae_override(pipe, diffusers_load_config, override_cls=diffusers.AutoencoderKLQwenImage, override_args={'low_cpu_mem_usage': False, 'ignore_mismatched_sizes': True})
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder
del transformer
del llm_adapter
diff --git a/pipelines/model_auraflow.py b/pipelines/model_auraflow.py
index 5d181ebce..597307c52 100644
--- a/pipelines/model_auraflow.py
+++ b/pipelines/model_auraflow.py
@@ -25,6 +25,8 @@ def load_auraflow(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder
del transformer
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py
index 4472a1a85..1e7f9ea8a 100644
--- a/pipelines/model_chroma.py
+++ b/pipelines/model_chroma.py
@@ -28,6 +28,9 @@ def load_chroma(checkpoint_info, diffusers_load_config=None):
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaPipeline
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaImg2ImgPipeline
diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["chroma"] = diffusers.ChromaInpaintPipeline
+
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder
del transformer
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_ernie.py b/pipelines/model_ernie.py
index c5ca50691..382b1fe80 100644
--- a/pipelines/model_ernie.py
+++ b/pipelines/model_ernie.py
@@ -40,6 +40,13 @@ def load_ernie_image(checkpoint_info, diffusers_load_config=None):
'use_pe': shared.opts.model_ernie_enable_pe,
}
+ from pipelines.ernie.ernie_image import ErnieImageImg2ImgPipeline, ErnieImageInpaintPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["ernieimage"] = diffusers.ErnieImagePipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["ernieimage"] = ErnieImageImg2ImgPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["ernieimage"] = ErnieImageInpaintPipeline
+
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del transformer
del text_encoder
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py
index 23bd40455..17686e95e 100644
--- a/pipelines/model_flux.py
+++ b/pipelines/model_flux.py
@@ -60,6 +60,8 @@ def load_flux(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
if os.environ.get('SD_REMOTE_T5', None) is not None:
from modules import sd_te_remote
log.warning('Remote-TE: applying patch')
diff --git a/pipelines/model_flux2.py b/pipelines/model_flux2.py
index dd6c364b0..a203730cb 100644
--- a/pipelines/model_flux2.py
+++ b/pipelines/model_flux2.py
@@ -31,6 +31,8 @@ def load_flux2(checkpoint_info, diffusers_load_config=None):
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux2"] = diffusers.Flux2Pipeline
diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux2"] = diffusers.Flux2Pipeline
+ generic.load_vae_override(pipe, diffusers_load_config)
+
from pipelines.flux import flux2_lora
flux2_lora.apply_patch()
diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py
index 31e539128..3a238805b 100644
--- a/pipelines/model_flux2_klein.py
+++ b/pipelines/model_flux2_klein.py
@@ -34,6 +34,8 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None):
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline
diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline
+ generic.load_vae_override(pipe, diffusers_load_config)
+
from pipelines.flux import flux2_lora
flux2_lora.apply_patch()
diff --git a/pipelines/model_google.py b/pipelines/model_google.py
index 91045fc53..40e52c5bd 100644
--- a/pipelines/model_google.py
+++ b/pipelines/model_google.py
@@ -28,8 +28,6 @@ aspect_ratios_buckets = {
def google_requirements():
from installer import install # , reload
install('google-genai==1.52.0')
- # install('pydantic==2.11.7', ignore=True, quiet=True)
- # reload('pydantic', '2.11.7')
def get_size_buckets(width: int, height: int) -> tuple[str, str]:
@@ -42,6 +40,7 @@ def get_size_buckets(width: int, height: int) -> tuple[str, str]:
class GoogleNanoBananaPipeline():
def __init__(self, model_name: str):
+ self.skip_processing = True
self.model = model_name
self.client = None
self.config = None
@@ -55,15 +54,18 @@ class GoogleNanoBananaPipeline():
contents=prompt,
)
- def img2img(self, prompt, image):
+ def img2img(self, prompt, images):
from google import genai # pylint: disable=no-name-in-module
- image_bytes = io.BytesIO()
- image.save(image_bytes, format='JPEG')
+ image_bytes_list = []
+ for image in images:
+ image_bytes = io.BytesIO()
+ image.save(image_bytes, format='JPEG')
+ image_bytes_list.append(genai.types.Part.from_bytes(data=image_bytes.getvalue(), mime_type='image/jpeg'))
return self.client.models.generate_content(
model=self.model,
config=self.config,
contents=[
- genai.types.Part.from_bytes(data=image_bytes.getvalue(), mime_type='image/jpeg'),
+ *image_bytes_list,
prompt,
],
)
@@ -108,7 +110,7 @@ class GoogleNanoBananaPipeline():
log.debug(f'Cloud: model="{self.model}" args={args_log}')
return args
- def __call__(self, prompt: list[str], width: int, height: int, image: Image.Image = None):
+ def __call__(self, prompt: list[str], width: int, height: int, images: list[Image.Image] = []):
from google import genai # pylint: disable=no-name-in-module
if self.client is None:
args = self.get_args()
@@ -125,13 +127,13 @@ class GoogleNanoBananaPipeline():
response_modalities=["IMAGE"],
image_config=image_config
)
- log.debug(f'Cloud: model="{self.model}" prompt="{prompt}" size={image_size} ar={aspect_ratio} image={image}')
+ log.debug(f'Cloud: model="{self.model}" prompt="{prompt}" size={image_size} ar={aspect_ratio} images={len(images) if images is not None else 0}')
# log.debug(f'Cloud: config={self.config}')
try:
t0 = time.time()
- if image is not None:
- response = self.img2img(prompt, image)
+ if images is not None and len(images) > 0:
+ response = self.img2img(prompt, images)
else:
response = self.txt2img(prompt)
t1 = time.time()
diff --git a/pipelines/model_hidream.py b/pipelines/model_hidream.py
index 1e05fbfef..7d1fcd5e4 100644
--- a/pipelines/model_hidream.py
+++ b/pipelines/model_hidream.py
@@ -31,6 +31,56 @@ def load_llama(diffusers_load_config=None):
return text_encoder_4, tokenizer_4
+def load_hidream_o1(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)
+
+ from pipelines.hidream.hidream_o1 import HiDreamO1Pipeline, HiDreamO1ImagePipeline
+ from pipelines.hidream.qwen3_vl_transformers import HiDreamO1Qwen3VLTransformer
+ from pipelines.hidream.scheduler_flashfloweuler import FlashFlowMatchEulerDiscreteScheduler
+
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, allow_quant=True)
+ log.debug(f'Load model: type=HiDreamO1 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+
+ o1_load_config = diffusers_load_config.copy()
+ o1_load_config['trust_remote_code'] = True
+
+ transformer = HiDreamO1Qwen3VLTransformer.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.hfcache_dir,
+ trust_remote_code=True,
+ **load_args,
+ **quant_args,
+ )
+ if shared.opts.diffusers_offload_mode != 'none' and transformer is not None:
+ sd_models.move_model(transformer, devices.cpu)
+
+ processor = transformers.AutoProcessor.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.hfcache_dir,
+ trust_remote_code=True,
+ )
+ pipe = HiDreamO1Pipeline(
+ transformer=transformer,
+ processor=processor,
+ tokenizer=processor.tokenizer,
+ scheduler=FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False),
+ )
+ pipe.task_args = {
+ 'output_type': 'pil',
+ }
+
+ del processor
+ del transformer
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-o1"] = HiDreamO1Pipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-o1"] = HiDreamO1ImagePipeline
+
+ devices.torch_gc()
+ return pipe
+
+
def load_hidream(checkpoint_info, diffusers_load_config=None):
if diffusers_load_config is None:
diffusers_load_config = {}
@@ -52,7 +102,7 @@ def load_hidream(checkpoint_info, diffusers_load_config=None):
if 'I1' in repo_id:
cls = diffusers.HiDreamImagePipeline
elif 'E1' in repo_id:
- from pipelines.hidream.pipeline_hidream_image_editing import HiDreamImageEditingPipeline
+ from pipelines.hidream.hidream_e1 import HiDreamImageEditingPipeline
cls = HiDreamImageEditingPipeline
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-e1"] = diffusers.HiDreamImagePipeline
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline
diff --git a/pipelines/model_joy.py b/pipelines/model_joy.py
new file mode 100644
index 000000000..5ed363315
--- /dev/null
+++ b/pipelines/model_joy.py
@@ -0,0 +1,50 @@
+import diffusers
+import transformers
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_joy(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, allow_quant=False)
+ log.debug(f'Load model: type=JoyImageEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+ transformer = generic.load_transformer(
+ repo_id,
+ cls_name=diffusers.JoyImageEditTransformer3DModel,
+ load_config=diffusers_load_config,
+ )
+ text_encoder = generic.load_text_encoder(
+ repo_id,
+ cls_name=transformers.Qwen3VLForConditionalGeneration,
+ load_config=diffusers_load_config,
+ )
+
+ pipe = diffusers.JoyImageEditPipeline.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ transformer=transformer,
+ text_encoder=text_encoder,
+ **load_args,
+ )
+ pipe.task_args = {
+ 'output_type': 'np',
+ }
+
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = diffusers.JoyImageEditPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['joy-image-edit'] = diffusers.JoyImageEditPipeline
+
+ generic.load_vae_override(pipe, diffusers_load_config)
+
+ del transformer
+ del text_encoder
+ sd_hijack_te.init_hijack(pipe)
+ sd_hijack_vae.init_hijack(pipe)
+
+ devices.torch_gc(force=True, reason='load')
+ return pipe
diff --git a/pipelines/model_lumina.py b/pipelines/model_lumina.py
index c55d8281b..e4fb57daa 100644
--- a/pipelines/model_lumina.py
+++ b/pipelines/model_lumina.py
@@ -18,6 +18,9 @@ def load_lumina(checkpoint_info, diffusers_load_config=None):
cache_dir = shared.opts.diffusers_dir,
**load_config,
)
+
+ generic.load_vae_override(pipe, diffusers_load_config)
+
sd_hijack_te.init_hijack(pipe)
devices.torch_gc(force=True, reason='load')
return pipe
@@ -47,6 +50,8 @@ def load_lumina2(checkpoint_info, diffusers_load_config=None):
**load_config,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del transformer
del text_encoder
sd_hijack_te.init_hijack(pipe)
@@ -61,38 +66,31 @@ def load_lumina_dimoo(checkpoint_info, diffusers_load_config=None):
repo_id = sd_models.path_to_repo(checkpoint_info)
sd_models.hf_auth_check(checkpoint_info)
- load_config, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+ load_config, quant_args = model_quant.get_dit_args(diffusers_load_config)
log.debug(f'Load model: type=LuminaDiMOO repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_config}')
- pipe_cls = getattr(diffusers, 'LuminaDiMOOPipeline', None)
- if pipe_cls is not None:
- pipe = pipe_cls.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- **load_config,
- )
- else:
- try:
- pipe = diffusers.DiffusionPipeline.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- trust_remote_code=True,
- **load_config,
- )
- except Exception as e:
- raise RuntimeError(f'Lumina-DiMOO is not available in installed diffusers={diffusers.__version__}. Please update diffusers to a version that includes LuminaDiMOOPipeline.') from e
+ from pipelines.lumina_dimmo.pipelines import LuminaDiMOOTextPipeline, LuminaDiMOOImagePipeline
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["luminadimoo"] = LuminaDiMOOTextPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["luminadimoo"] = LuminaDiMOOImagePipeline
+ # Force slow tokenizer path for Lumina-DiMOO to avoid fast-tokenizer conversion failures.
+ tokenizer = transformers.AutoTokenizer.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ trust_remote_code=True,
+ use_fast=False,
+ )
+
+ pipe = LuminaDiMOOTextPipeline.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ tokenizer=tokenizer,
+ **load_config,
+ **quant_args,
+ )
+ pipe.skip_processing = True
+ pipe.task_args = {'output_type': 'np'}
+
+ del tokenizer
devices.torch_gc(force=True, reason='load')
return pipe
-
-""" Reference
- "AlphaVLLM Lumina DiMOO": {
- "path": "Alpha-VLLM/Lumina-DiMOO",
- "desc": "Lumina-DiMOO is an omni diffusion large language model for multimodal generation and understanding with text-to-image, image editing and understanding capabilities.",
- "preview": "Alpha-VLLM--Lumina-DiMOO.jpg",
- "skip": true,
- "extras": "sampler: Default",
- "size": 0,
- "date": "2025 September"
- },
-"""
diff --git a/pipelines/model_nucleus.py b/pipelines/model_nucleus.py
index 1dd4ca89b..b7087fa59 100644
--- a/pipelines/model_nucleus.py
+++ b/pipelines/model_nucleus.py
@@ -42,6 +42,8 @@ def load_nucleus(checkpoint_info, diffusers_load_config=None):
'output_type': 'np',
}
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del transformer
del text_encoder
del processor
diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py
index 2f91ac895..1ced8b659 100644
--- a/pipelines/model_pixart.py
+++ b/pipelines/model_pixart.py
@@ -35,6 +35,8 @@ def load_pixart(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder
del transformer
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py
index 36aa99888..52c9fd251 100644
--- a/pipelines/model_qwen.py
+++ b/pipelines/model_qwen.py
@@ -88,6 +88,8 @@ def load_qwen(checkpoint_info, diffusers_load_config=None):
pipe.task_args['layers'] = shared.opts.model_qwen_layers
pipe.task_args['resolution'] = 640
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder
del transformer
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_sd3.py b/pipelines/model_sd3.py
index c19cc1c11..5717dbc38 100644
--- a/pipelines/model_sd3.py
+++ b/pipelines/model_sd3.py
@@ -32,6 +32,8 @@ def load_sd3(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del text_encoder_3
del transformer
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/model_step1x_edit.py b/pipelines/model_step1x_edit.py
index 38c5f4284..aba2bd5ab 100644
--- a/pipelines/model_step1x_edit.py
+++ b/pipelines/model_step1x_edit.py
@@ -1,11 +1,14 @@
import transformers
import diffusers
-from modules import shared, devices, sd_models, model_quant, sd_hijack_te
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
from modules.logger import log
from pipelines import generic
def load_step1x_edit(checkpoint_info, diffusers_load_config=None):
+ from pipelines.step1x.pipeline_step1x_edit import Step1XEditPipeline
+ from pipelines.step1x.transformer_step1x_edit import Step1XEditTransformer2DModel
+
if diffusers_load_config is None:
diffusers_load_config = {}
repo_id = sd_models.path_to_repo(checkpoint_info)
@@ -14,53 +17,30 @@ def load_step1x_edit(checkpoint_info, diffusers_load_config=None):
load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
log.debug(f'Load model: type=Step1XEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
- # Load text encoder (Qwen2.5-VL - available in transformers)
- text_encoder = generic.load_text_encoder(repo_id,
- cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config
- )
+ diffusers.Step1XEditPipeline = Step1XEditPipeline
+ diffusers.Step1XEditTransformer2DModel = Step1XEditTransformer2DModel
- # Load processor for Qwen2.5-VL
- processor = transformers.Qwen2_5_VLProcessor.from_pretrained(repo_id,
- cache_dir=shared.opts.hfcache_dir,
- subfolder='processor'
- )
+ text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config)
+ processor = transformers.Qwen2_5_VLProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir, subfolder='processor')
+ transformer = generic.load_transformer(repo_id, cls_name=Step1XEditTransformer2DModel, load_config=diffusers_load_config)
- # Step1XEditPipeline and Step1XEditTransformer2DModel are custom classes not in current diffusers
- # Try direct pipeline class first, fall back to trust_remote_code
- pipe_cls = getattr(diffusers, 'Step1XEditPipeline', None)
- if pipe_cls is not None:
- pipe = pipe_cls.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- text_encoder=text_encoder,
- processor=processor,
- **load_args,
- )
- else:
- pipe = diffusers.DiffusionPipeline.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- text_encoder=text_encoder,
- processor=processor,
- trust_remote_code=True,
- **load_args,
- )
+ pipe = Step1XEditPipeline.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ transformer=transformer,
+ text_encoder=text_encoder,
+ processor=processor,
+ **load_args,
+ )
+ pipe.task_args = {
+ 'output_type': 'pil', # step1x is buggy with np
+ }
del text_encoder
del processor
+ del transformer
sd_hijack_te.init_hijack(pipe)
+ sd_hijack_vae.init_hijack(pipe)
devices.torch_gc(force=True, reason='load')
return pipe
-
-""" Reference
- "StepFun Step1X-Edit v1.2": {
- "path": "stepfun-ai/Step1X-Edit-v1p2",
- "preview": "stepfun-ai--Step1X-Edit-v1p2.jpg",
- "desc": "Multimodal image editing model using Step1X transformer architecture with Qwen2.5-VL text encoding, trained with Flow Matching scheduler for high-quality in-context image edits and refinements.",
- "skip": true,
- "extras": "sampler: Default",
- "size": 8.58,
- "date": "2025 October"
- },
-"""
diff --git a/pipelines/model_ultraflux.py b/pipelines/model_ultraflux.py
new file mode 100644
index 000000000..8146febc1
--- /dev/null
+++ b/pipelines/model_ultraflux.py
@@ -0,0 +1,55 @@
+import diffusers
+import transformers
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_ultraflux(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, allow_quant=False)
+ log.debug(f'Load model: type=UltraFlux repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+ from pipelines.ultraflux.pipeline_flux import UltraFluxPipeline
+ from pipelines.ultraflux.transformer_flux import FluxTransformer2DModel
+ from pipelines.ultraflux.autoencoder_kl import AutoencoderUltraFluxKL
+
+ transformer = generic.load_transformer(repo_id, cls_name=FluxTransformer2DModel, load_config=diffusers_load_config)
+ text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder='text_encoder_2')
+ vae = AutoencoderUltraFluxKL.from_pretrained(
+ repo_id,
+ subfolder='vae',
+ cache_dir=shared.opts.diffusers_dir,
+ torch_dtype=devices.dtype,
+ )
+ pipe = UltraFluxPipeline.from_pretrained(
+ repo_id,
+ transformer=transformer,
+ text_encoder_2=text_encoder_2,
+ vae=vae,
+ cache_dir=shared.opts.diffusers_dir,
+ **load_args,
+ )
+ pipe.task_args = {
+ 'output_type': 'np',
+ }
+ if hasattr(pipe, 'scheduler') and hasattr(pipe.scheduler, 'config'):
+ if hasattr(pipe.scheduler.config, 'use_dynamic_shifting'):
+ pipe.scheduler.config.use_dynamic_shifting = False
+ if hasattr(pipe.scheduler.config, 'time_shift'):
+ pipe.scheduler.config.time_shift = 4
+
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['ultraflux'] = UltraFluxPipeline
+
+ del text_encoder_2
+ del transformer
+ del vae
+ sd_hijack_te.init_hijack(pipe)
+ sd_hijack_vae.init_hijack(pipe)
+
+ devices.torch_gc(force=True, reason='load')
+ return pipe
diff --git a/pipelines/model_vibe.py b/pipelines/model_vibe.py
index fbaec7ede..ecac08867 100644
--- a/pipelines/model_vibe.py
+++ b/pipelines/model_vibe.py
@@ -1,3 +1,4 @@
+import sys
import diffusers
import transformers
from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
@@ -14,58 +15,46 @@ def load_vibe(checkpoint_info, diffusers_load_config=None):
load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
log.debug(f'Load model: type=VIBE repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
- pipe_cls = getattr(diffusers, 'VIBESanaEditingPipeline', None)
- transformer_cls = getattr(diffusers, 'VIBESanaEditingModel', None)
- text_encoder_cls = getattr(transformers, 'Qwen3VLForConditionalGeneration', None)
- processor_cls = getattr(transformers, 'Qwen3VLProcessor', None)
+ from pipelines.vibe import VIBESanaEditingModel, VIBESanaEditingPipeline, VIBESanaImagePipeline
+ diffusers.VIBESanaEditingPipeline = VIBESanaEditingPipeline
+ diffusers.VIBESanaEditingModel = VIBESanaEditingModel
- if pipe_cls is not None and transformer_cls is not None and text_encoder_cls is not None and processor_cls is not None:
- transformer = generic.load_transformer(
- repo_id,
- cls_name=transformer_cls,
- load_config=diffusers_load_config,
- allow_quant=False,
- )
- text_encoder = generic.load_text_encoder(
- repo_id,
- cls_name=text_encoder_cls,
- load_config=diffusers_load_config,
- allow_quant=False,
- allow_shared=False,
- )
- processor = processor_cls.from_pretrained(repo_id, subfolder='tokenizer', cache_dir=shared.opts.hfcache_dir)
+ sys.modules['vibe.transformer.vibe_sana_editing'] = diffusers # monkey patch since hf model_index.json points to custom class path
- pipe = pipe_cls.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- transformer=transformer,
- text_encoder=text_encoder,
- tokenizer=processor,
- **load_args,
- )
+ transformer = generic.load_transformer(
+ repo_id,
+ cls_name=VIBESanaEditingModel,
+ load_config=diffusers_load_config,
+ allow_quant=False,
+ )
+ text_encoder = generic.load_text_encoder(
+ repo_id,
+ cls_name=transformers.Qwen3VLForConditionalGeneration,
+ load_config=diffusers_load_config,
+ allow_quant=False,
+ allow_shared=False,
+ )
+ processor = transformers.Qwen3VLProcessor.from_pretrained(
+ repo_id,
+ subfolder='tokenizer',
+ cache_dir=shared.opts.hfcache_dir,
+ )
- del transformer
- del text_encoder
- del processor
- else:
- try:
- from installer import install, installed
- if not installed('vibe', quiet=True):
- install('git+https://github.com/ai-forever/VIBE', 'vibe')
- import vibe # pylint: disable=unused-import
- except Exception as e:
- raise RuntimeError('VIBE requires either native diffusers VIBESana classes or `vibe` package') from e
+ pipe = VIBESanaEditingPipeline.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ transformer=transformer,
+ text_encoder=text_encoder,
+ tokenizer=processor,
+ **load_args,
+ )
- pipe = diffusers.DiffusionPipeline.from_pretrained(
- repo_id,
- cache_dir=shared.opts.diffusers_dir,
- trust_remote_code=True,
- **load_args,
- )
+ del transformer
+ del text_encoder
+ del processor
- diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['vibe-sana'] = pipe.__class__
- diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['vibe-sana'] = pipe.__class__
- diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING['vibe-sana'] = pipe.__class__
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['vibe-sana'] = VIBESanaEditingPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['vibe-sana'] = VIBESanaImagePipeline
pipe.task_args = {
'output_type': 'np',
@@ -76,15 +65,3 @@ def load_vibe(checkpoint_info, diffusers_load_config=None):
devices.torch_gc(force=True, reason='load')
return pipe
-
-""" Reference
- "VIBE Image Edit": {
- "path": "iitolstykh/VIBE-Image-Edit",
- "preview": "iitolstykh--VIBE-Image-Edit.jpg",
- "desc": "VIBE is an open-source text-guided image editing model combining Sana1.5-1.6B diffusion backbone with Qwen3-VL multimodal conditioning for fast, instruction-based edits.",
- "skip": true,
- "extras": "sampler: Default, cfg_scale: 4.5, image_guidance_scale: 1.2, steps: 20",
- "size": 9.72,
- "date": "2025 December"
- },
-"""
diff --git a/pipelines/model_z_image.py b/pipelines/model_z_image.py
index 3081bb595..e619de619 100644
--- a/pipelines/model_z_image.py
+++ b/pipelines/model_z_image.py
@@ -52,6 +52,8 @@ def load_z_image(checkpoint_info, diffusers_load_config=None):
**load_args,
)
+ generic.load_vae_override(pipe, diffusers_load_config)
+
del transformer
del text_encoder
sd_hijack_te.init_hijack(pipe)
diff --git a/pipelines/omnigen2/image_processor.py b/pipelines/omnigen2/image_processor.py
index 00b1fb3ed..ea908e7db 100644
--- a/pipelines/omnigen2/image_processor.py
+++ b/pipelines/omnigen2/image_processor.py
@@ -192,14 +192,14 @@ class OmniGen2ImageProcessor(VaeImageProcessor):
warnings.warn(
"Passing `image` as a list of 4d np.ndarray is deprecated."
"Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
- FutureWarning,
+ FutureWarning, stacklevel=2,
)
image = np.concatenate(image, axis=0)
if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
warnings.warn(
"Passing `image` as a list of 4d torch.Tensor is deprecated."
"Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
- FutureWarning,
+ FutureWarning, stacklevel=2,
)
image = torch.cat(image, axis=0)
@@ -253,7 +253,7 @@ class OmniGen2ImageProcessor(VaeImageProcessor):
warnings.warn(
"Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
- FutureWarning,
+ FutureWarning, stacklevel=2,
)
do_normalize = False
if do_normalize:
diff --git a/pipelines/omnigen2/models/transformers/repo.py b/pipelines/omnigen2/models/transformers/repo.py
index 8f7c47566..133d800ab 100644
--- a/pipelines/omnigen2/models/transformers/repo.py
+++ b/pipelines/omnigen2/models/transformers/repo.py
@@ -23,7 +23,7 @@ class OmniGen2RotaryPosEmbed(nn.Module):
theta: int) -> List[torch.Tensor]:
freqs_cis = []
freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
- for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
+ for _i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
freqs_cis.append(emb)
return freqs_cis
@@ -114,7 +114,7 @@ class OmniGen2RotaryPosEmbed(nn.Module):
batch_size, max_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
)
- for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)):
+ for i, (cap_seq_len, ref_img_len, img_len, _) in enumerate(zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)):
cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len]
ref_img_freqs_cis[i, :sum(ref_img_len)] = freqs_cis[i, cap_seq_len:cap_seq_len + sum(ref_img_len)]
img_freqs_cis[i, :img_len] = freqs_cis[i, cap_seq_len + sum(ref_img_len):cap_seq_len + sum(ref_img_len) + img_len]
diff --git a/pipelines/omnigen2/models/transformers/transformer_omnigen2.py b/pipelines/omnigen2/models/transformers/transformer_omnigen2.py
index fe65b3170..a530245ff 100644
--- a/pipelines/omnigen2/models/transformers/transformer_omnigen2.py
+++ b/pipelines/omnigen2/models/transformers/transformer_omnigen2.py
@@ -444,7 +444,7 @@ class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From
if ref_img_sizes[i] is not None:
imgs = []
for ref_img in ref_image_hidden_states[i]:
- C, H, W = ref_img.size()
+ _C, _H, _W = ref_img.size()
ref_img = rearrange(ref_img, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=p, p2=p)
imgs.append(ref_img)
@@ -457,7 +457,7 @@ class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From
flat_hidden_states = []
for i in range(batch_size):
img = hidden_states[i]
- C, H, W = img.size()
+ _C, _H, _W = img.size()
img = rearrange(img, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=p, p2=p)
flat_hidden_states.append(img)
@@ -518,7 +518,7 @@ class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From
if is_hidden_states_tensor:
assert hidden_states.ndim == 4
- hidden_states = [_hidden_states for _hidden_states in hidden_states]
+ hidden_states = list(hidden_states)
device = hidden_states[0].device
@@ -580,7 +580,7 @@ class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From
hidden_states = joint_hidden_states
- for layer_idx, layer in enumerate(self.layers):
+ for _layer_idx, layer in enumerate(self.layers):
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states = self._gradient_checkpointing_func(
layer, hidden_states, attention_mask, rotary_emb, temb
diff --git a/pipelines/omnigen2/pipeline_omnigen2.py b/pipelines/omnigen2/pipeline_omnigen2.py
index d2e13ba18..226d8ec99 100644
--- a/pipelines/omnigen2/pipeline_omnigen2.py
+++ b/pipelines/omnigen2/pipeline_omnigen2.py
@@ -243,10 +243,10 @@ class OmniGen2Pipeline(DiffusionPipeline):
if batch_size == 1:
images = [images]
latents = []
- for i, img in enumerate(images):
+ for _i, img in enumerate(images):
if img is not None and len(img) > 0:
ref_latents = []
- for j, img_j in enumerate(img):
+ for _j, img_j in enumerate(img):
img_j = self.image_processor.preprocess(img_j, max_pixels=max_pixels, max_side_length=max_side_length)
ref_latents.append(self.encode_vae(img_j.to(device=device)).squeeze(0))
else:
@@ -471,7 +471,7 @@ class OmniGen2Pipeline(DiffusionPipeline):
image_guidance_scale: float = 1.0,
cfg_range: Tuple[float, float] = (0.0, 1.0),
attention_kwargs: Optional[Dict[str, Any]] = None,
- timesteps: List[int] = None,
+ timesteps: List[int] | None = None,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
@@ -610,7 +610,7 @@ class OmniGen2Pipeline(DiffusionPipeline):
verbose,
step_func=None
):
- batch_size = latents.shape[0]
+ latents.shape[0]
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
@@ -701,7 +701,7 @@ class OmniGen2Pipeline(DiffusionPipeline):
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep = t.expand(latents.shape[0]).to(latents.dtype)
- batch_size, num_channels_latents, height, width = latents.shape
+ _batch_size, _num_channels_latents, _height, _width = latents.shape
optional_kwargs = {}
if 'ref_image_hidden_states' in set(inspect.signature(self.transformer.forward).parameters.keys()):
diff --git a/pipelines/segmoe/segmoe_model.py b/pipelines/segmoe/segmoe_model.py
index a542c1fbb..a84ad30e7 100644
--- a/pipelines/segmoe/segmoe_model.py
+++ b/pipelines/segmoe/segmoe_model.py
@@ -166,7 +166,7 @@ class SegMoEPipeline:
if not os.path.isfile("base/model.safetensors"):
os.system(
"wget -O "
- + "base/model.safetensors"
+ "base/model.safetensors"
+ self.config["base_model"]
+ " --content-disposition"
)
@@ -221,8 +221,8 @@ class SegMoEPipeline:
if not os.path.isfile(f"expert_{i}/model.safetensors"):
os.system(
f"wget {exp['source_model']} -O "
- + f"expert_{i}/model.safetensors"
- + " --content-disposition"
+ f"expert_{i}/model.safetensors"
+ " --content-disposition"
)
exp["source_model"] = f"expert_{i}/model.safetensors"
expert = DiffusionPipeline.from_single_file(
@@ -267,8 +267,8 @@ class SegMoEPipeline:
):
os.system(
f"wget {lora['source_model']} -O "
- + f"expert_{i}/lora_{j}/pytorch_lora_weights.safetensors"
- + " --content-disposition"
+ f"expert_{i}/lora_{j}/pytorch_lora_weights.safetensors"
+ " --content-disposition"
)
lora["source_model"] = f"expert_{j}/lora_{j}"
expert.load_lora_weights(lora["source_model"])
@@ -299,8 +299,8 @@ class SegMoEPipeline:
):
os.system(
f"wget {lora['source_model']} -O "
- + f"lora_{i}/pytorch_lora_weights.safetensors"
- + " --content-disposition"
+ f"lora_{i}/pytorch_lora_weights.safetensors"
+ " --content-disposition"
)
lora["source_model"] = f"lora_{i}"
self.pipe.load_lora_weights(lora["source_model"])
@@ -338,8 +338,8 @@ class SegMoEPipeline:
):
os.system(
f"wget {lora['source_model']} -O "
- + f"lora_{i}/pytorch_lora_weights.safetensors"
- + " --content-disposition"
+ f"lora_{i}/pytorch_lora_weights.safetensors"
+ " --content-disposition"
)
lora["source_model"] = f"lora_{i}"
experts[j[i]].load_lora_weights(lora["source_model"])
diff --git a/pipelines/step1x/__init__.py b/pipelines/step1x/__init__.py
new file mode 100644
index 000000000..f7983b96f
--- /dev/null
+++ b/pipelines/step1x/__init__.py
@@ -0,0 +1,9 @@
+from pipelines.step1x.pipeline_output import Step1XEditPipelineOutput
+from pipelines.step1x.pipeline_step1x_edit import Step1XEditPipeline
+from pipelines.step1x.transformer_step1x_edit import Step1XEditTransformer2DModel
+
+__all__ = [
+ "Step1XEditPipeline",
+ "Step1XEditPipelineOutput",
+ "Step1XEditTransformer2DModel",
+]
diff --git a/pipelines/step1x/pipeline_output.py b/pipelines/step1x/pipeline_output.py
new file mode 100644
index 000000000..289f3aea1
--- /dev/null
+++ b/pipelines/step1x/pipeline_output.py
@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class Step1XEditPipelineOutput(BaseOutput):
+ """
+ Output class for Step1X-Edit pipelines.
+
+ Args:
+ images (`List[PIL.Image.Image]` or `np.ndarray`)
+ List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+ """
+
+ images: Union[List[PIL.Image.Image], np.ndarray]
diff --git a/pipelines/step1x/pipeline_step1x_edit.py b/pipelines/step1x/pipeline_step1x_edit.py
new file mode 100644
index 000000000..7c4f0e228
--- /dev/null
+++ b/pipelines/step1x/pipeline_step1x_edit.py
@@ -0,0 +1,1136 @@
+# Copyright 2025 Step1X-Edit Team 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.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+from PIL import Image
+import torch
+import math
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.models import AutoencoderKL
+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.pipeline_utils import DiffusionPipeline
+from pipelines.step1x.transformer_step1x_edit import Step1XEditTransformer2DModel
+from pipelines.step1x.pipeline_output import Step1XEditPipelineOutput
+
+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 Step1XEditPipeline
+ >>> from diffusers.utils import load_image
+
+ >>> pipe = Step1XEditPipeline.from_pretrained("stepfun-ai/Step1X-Edit-v1p1-diffusers", torch_dtype=torch.bfloat16)
+ >>> pipe.to("cuda")
+ >>> image = load_image(
+ ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
+ ... ).convert("RGB")
+ >>> prompt = "Make Pikachu hold a sign that says 'Step1X-Edit is awesome', yarn art style, detailed, vibrant colors"
+
+ >>> image = pipe(
+ image=image,
+ prompt=prompt,
+ num_inference_steps=28,
+ true_cfg_scale=6.0,
+ generator=torch.Generator().manual_seed(42),
+ ).images[0]
+ >>> image.save("output.png")
+ ```
+"""
+
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class Step1XEditPipeline(DiffusionPipeline):
+ r"""
+ The Step1X-Edit pipeline for image-to-image and text-to-image generation.
+
+ Reference: https://arxiv.org/abs/2504.17761
+
+ Args:
+ transformer ([`Step1XEditTransformer2DModel`]):
+ Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+ scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+ vae ([`AutoencoderKL`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+ text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
+ [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct)
+ processor (`Qwen2_5_VLProcessor`):
+ [Qwen2_5_VLProcessor](https://huggingface.co/docs/transformers/v4.53.3/en/model_doc/qwen2_5_vl#transformers.Qwen2_5_VLProcessor).
+ """
+
+ model_cpu_offload_seq = "text_encoder->transformer->vae"
+ _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+ def __init__(
+ self,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ vae: AutoencoderKL,
+ text_encoder: Qwen2_5_VLForConditionalGeneration,
+ processor: Qwen2_5_VLProcessor,
+ transformer: Step1XEditTransformer2DModel,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ processor=processor,
+ transformer=transformer,
+ scheduler=scheduler,
+ )
+ self.image_encoder=None
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+ # Step1X-Edit latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+ # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+ self.latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+ self.max_token_length = 640
+ self.default_sample_size = 128
+ self.QWEN25VL_PREFIX = '''Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:
+- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.
+- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n
+Here are examples of how to transform or refine prompts:
+- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.
+- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n
+Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:
+User Prompt:'''
+
+ def _split_string(self, s):
+ s = s.replace("'", '"').replace("“", '"').replace("”", '"') # use english quotes
+ result = []
+ in_quotes = False
+ temp = ""
+
+ for idx,char in enumerate(s):
+ if char == '"' and idx>155: # system token
+ temp += char
+ if not in_quotes:
+ result.append(temp)
+ temp = ""
+
+ in_quotes = not in_quotes
+ continue
+ if in_quotes:
+ if char.isspace():
+ pass # have space token
+
+ result.append("“" + char + "”")
+ else:
+ temp += char
+
+ if temp:
+ result.append(temp)
+
+ return result
+
+ def _get_qwenvl_embeds(
+ self,
+ prompt: Union[str, List[str]],
+ ref_image: Optional[torch.Tensor],
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = torch.bfloat16,
+ ):
+ text_list = prompt
+ embs = torch.zeros(
+ len(text_list),
+ self.max_token_length,
+ self.text_encoder.config.text_config.hidden_size,
+ dtype=dtype,
+ device=device,
+ )
+ masks = torch.zeros(
+ len(text_list),
+ self.max_token_length,
+ dtype=torch.long,
+ device=device,
+ )
+
+ if isinstance(ref_image, list):
+ image_list = ref_image
+ else:
+ image_list = [ref_image] * len(text_list)
+
+ if len(image_list) == 1 and len(text_list) > 1:
+ image_list = image_list * len(text_list)
+ if len(image_list) != len(text_list):
+ raise ValueError(
+ f"Mismatch between prompts ({len(text_list)}) and reference images ({len(image_list)})."
+ )
+
+ for idx, (txt, imgs) in enumerate(zip(text_list, image_list)):
+
+ messages = [
+ {
+ "role": "user",
+ "content": []
+ }
+ ]
+
+ messages[0]["content"].append({"type": "text", "text": f"{self.QWEN25VL_PREFIX}"})
+ messages[0]['content'].append({"type": "image", "image": imgs})
+ messages[0]["content"].append({"type": "text", "text": f"{txt}"})
+
+ # Preparation for inference
+ text = self.processor.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True, add_vision_id=True
+ )
+ imgs = imgs.convert("RGB")
+ min_pixels = 4 * 28 * 28
+ max_pixels = 16384 * 28 * 28
+ width, height = imgs.size
+ h_bar = max(28, round(height / 28) * 28)
+ w_bar = max(28, round(width / 28) * 28)
+ if h_bar * w_bar > max_pixels:
+ beta = math.sqrt((height * width) / max_pixels)
+ h_bar = math.floor(height / beta / 28) * 28
+ w_bar = math.floor(width / beta / 28) * 28
+ elif h_bar * w_bar < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ h_bar = math.ceil(height * beta / 28) * 28
+ w_bar = math.ceil(width * beta / 28) * 28
+ image_inputs = [imgs.resize((w_bar, h_bar))]
+
+ inputs = self.processor(
+ text=[text],
+ images=image_inputs,
+ padding=True,
+ return_tensors="pt",
+ )
+
+ old_inputs_ids = inputs.input_ids
+ text_split_list = self._split_string(text)
+
+ token_list = []
+ for text_each in text_split_list:
+ txt_inputs = self.processor(
+ text=text_each,
+ images=None,
+ videos=None,
+ padding=True,
+ return_tensors="pt",
+ )
+ token_each=txt_inputs.input_ids
+ if token_each[0][0] == 2073 and token_each[0][-1] == 854:
+ token_each = token_each[:,1:-1]
+ token_list.append(token_each)
+ else:
+ token_list.append(token_each)
+
+ new_txt_ids = torch.cat(token_list, dim=1).to(device)
+
+ new_txt_ids = new_txt_ids.to(old_inputs_ids.device)
+ idx1 = (old_inputs_ids == 151653).nonzero(as_tuple=True)[1][0]
+ idx2 = (new_txt_ids == 151653).nonzero(as_tuple=True)[1][0]
+ inputs.input_ids = torch.cat([old_inputs_ids[0, :idx1], new_txt_ids[0, idx2:]], dim=0).unsqueeze(0).to(device)
+ inputs.attention_mask = (inputs.input_ids > 0).long().to(device)
+ outputs = self.text_encoder(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, pixel_values=inputs.pixel_values.to(device), image_grid_thw=inputs.image_grid_thw.to(device), output_hidden_states=True)
+
+ emb = outputs['hidden_states'][-1]
+ embs[idx,:min(self.max_token_length,emb.shape[1]-217)] = emb[0,217:][:self.max_token_length]
+ masks[idx, :min(self.max_token_length, emb.shape[1] - 217)] = torch.ones((min(self.max_token_length, emb.shape[1] - 217)), dtype=torch.long, device=device)
+
+ return embs, masks
+
+ def encode_prompt(
+ self,
+ ref_image: Optional[torch.Tensor],
+ prompt: Union[str, List[str]],
+ device: Optional[torch.device] = None,
+ num_images_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds_mask: Optional[torch.Tensor] = None,
+ ):
+ r"""
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ """
+ device = device or self._execution_device
+
+ ref_image = [ref_image] if isinstance(prompt, str) else ref_image # change
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds, prompt_embeds_mask = self._get_qwenvl_embeds(prompt, ref_image, device)
+
+ _, 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)
+ prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+ text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device)
+
+ return prompt_embeds, prompt_embeds_mask, text_ids
+
+ def encode_image(
+ self,
+ image: Optional[torch.Tensor],
+ width: Optional[int] = None,
+ height: Optional[int] = None,
+ device: Optional[torch.device] = None,
+ num_images_per_prompt: int = 1,
+ ):
+
+ if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
+ img_info = image.size
+ width, height = img_info
+ aspect_ratio = width / height
+
+ if width > height:
+ width_new = math.ceil(math.sqrt(1024 * 1024 * aspect_ratio))
+ height_new = math.ceil(width_new / aspect_ratio)
+ else:
+ height_new = math.ceil(math.sqrt(1024 * 1024 / aspect_ratio))
+ width_new = math.ceil(height_new * aspect_ratio)
+
+ multiple_of = self.vae_scale_factor * 2
+ height_new = height_new // multiple_of * multiple_of
+ width_new = width_new // multiple_of * multiple_of
+
+ if height != height_new or width != width_new:
+ logger.warning(
+ f"Generation `height` and `width` have been adjusted to {height_new} and {width_new} to fit the model requirements."
+ )
+ height, width = height_new, width_new
+ ref_image = self.image_processor.resize(image, height, width)
+ image = self.image_processor.preprocess(ref_image, height, width).contiguous()
+ else:
+ width = width if width is not None else 1024
+ height = height if height is not None else 1024
+ img_info = (width, height)
+ # Keep t2i reference image size aligned with requested output size.
+ ref_image = torch.zeros(3, height, width).unsqueeze(0).to(device)
+ ref_image = self.image_processor.pt_to_numpy(ref_image)
+ ref_image = self.image_processor.numpy_to_pil(ref_image)[0]
+ image = None
+
+ return image, ref_image, img_info, width, height
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.check_inputs
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_embeds_mask=None,
+ negative_prompt_embeds_mask=None,
+ callback_on_step_end_tensor_inputs=None,
+ ):
+ if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+ logger.warning(
+ f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+ )
+
+ 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 {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and prompt_embeds_mask is None:
+ raise ValueError(
+ "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`."
+ )
+ if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+ raise ValueError(
+ "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`."
+ )
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids
+ def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+ latent_image_ids = torch.zeros(height, width, 3)
+ latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+ latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+ latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+ latent_image_ids = latent_image_ids.reshape(
+ latent_image_id_height * latent_image_id_width, latent_image_id_channels
+ )
+
+ return latent_image_ids.to(device=device, dtype=dtype)
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents
+ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+ latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+ latents = latents.permute(0, 2, 4, 1, 3, 5)
+ latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+ return latents
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents
+ def _unpack_latents(latents, height, width, vae_scale_factor):
+ batch_size, _num_patches, channels = latents.shape
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = 2 * (int(height) // (vae_scale_factor * 2))
+ width = 2 * (int(width) // (vae_scale_factor * 2))
+
+ latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+ latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+ return latents
+
+ @staticmethod
+ def _output_process_image(image, image_size):
+ def _to_pil(img):
+ arr = np.asarray(img)
+
+ # Handle CHW tensors by converting to HWC for PIL.
+ if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+ arr = np.transpose(arr, (1, 2, 0))
+
+ if np.issubdtype(arr.dtype, np.floating):
+ arr = np.nan_to_num(arr)
+ if arr.size > 0:
+ arr_min = float(arr.min())
+ arr_max = float(arr.max())
+ if arr_min >= -1.0 and arr_max <= 1.0:
+ if arr_min < 0.0:
+ arr = (arr + 1.0) * 0.5
+ arr = arr * 255.0
+ arr = np.clip(arr, 0.0, 255.0).astype(np.uint8)
+ elif arr.dtype != np.uint8:
+ arr = np.clip(arr, 0, 255).astype(np.uint8)
+
+ return Image.fromarray(arr)
+
+ if isinstance(image, np.ndarray):
+ images = image if image.ndim == 4 else [image]
+ resized = [np.array(_to_pil(img).resize(image_size)) for img in images]
+ output = np.stack(resized, axis=0)
+ return output if image.ndim == 4 else output[0]
+ if len(image) > 0 and isinstance(image[0], np.ndarray):
+ resized = [np.array(_to_pil(img).resize(image_size)) for img in image]
+ return np.stack(resized, axis=0)
+ return [img.resize(image_size) for img in image]
+
+ @staticmethod
+ def process_diff_norm(diff_norm, k):
+ pow_result = torch.pow(diff_norm, k)
+
+ result = torch.where(
+ diff_norm > 1.0,
+ pow_result,
+ torch.where(diff_norm < 1.0, torch.ones_like(diff_norm), diff_norm),
+ )
+ return result
+
+ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+ image_latents = None
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="sample")
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="sample")
+
+ image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+ return image_latents
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing
+ def enable_vae_slicing(self):
+ r"""
+ Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+ compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+ """
+ self.vae.enable_slicing()
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing
+ def disable_vae_slicing(self):
+ r"""
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+ computing decoding in one step.
+ """
+ self.vae.disable_slicing()
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling
+ def enable_vae_tiling(self):
+ r"""
+ Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+ compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+ processing larger images.
+ """
+ self.vae.enable_tiling()
+
+ # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling
+ def disable_vae_tiling(self):
+ r"""
+ Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+ computing decoding in one step.
+ """
+ self.vae.disable_tiling()
+
+ def prepare_latents(
+ self,
+ image: Optional[torch.Tensor],
+ batch_size: int,
+ num_channels_latents: int,
+ height: int,
+ width: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ ):
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = 2 * (int(height) // (self.vae_scale_factor * 2))
+ width = 2 * (int(width) // (self.vae_scale_factor * 2))
+ shape = (batch_size, num_channels_latents, height, width)
+
+ image_latents = image_ids = None
+ if image is not None:
+ image = image.to(device=device, dtype=dtype)
+
+ if image.shape[1] != self.latent_channels:
+ image_latents = self._encode_vae_image(image=image, generator=generator)
+ else:
+ image_latents = image
+ if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+ # expand init_latents for batch_size
+ additional_image_per_prompt = batch_size // image_latents.shape[0]
+ image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+ elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+ raise ValueError(
+ f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+ )
+ else:
+ image_latents = torch.cat([image_latents], dim=0)
+
+ image_latent_height, image_latent_width = image_latents.shape[2:]
+ image_latents = self._pack_latents(
+ image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
+ )
+ image_ids = self._prepare_latent_image_ids(
+ batch_size, image_latent_height // 2, image_latent_width // 2, device, torch.float32 # change
+ # batch_size, image_latent_height // 2, image_latent_width // 2, device, dtype
+ )
+ # image ids are the same as latent ids with the first dimension set to 1 instead of 0
+ image_ids[..., 0] = 1
+ image_ids[..., 1] = image_ids[..., 1] + 1
+ image_ids[..., 2] = image_ids[..., 2] + 1
+ latent_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) # change
+ latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+ else:
+ latents = latents.to(device=device, dtype=dtype)
+
+ return latents, image_latents, latent_ids, image_ids
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def joint_attention_kwargs(self):
+ return self._joint_attention_kwargs
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def current_timestep(self):
+ return self._current_timestep
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ image: Optional[PipelineImageInput] = None,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ true_cfg_scale: float = 6.0,
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ num_inference_steps: int = 28,
+ sigmas: Optional[List[float]] = None,
+ guidance_scale: float = 6.0,
+ num_images_per_prompt: int = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds_mask: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+ ip_adapter_image: Optional[PipelineImageInput] = None,
+ ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+ negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+ negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ timesteps_truncate: float = 0.93,
+ process_norm_power: float = 0.4
+ ):
+ r"""
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
+ numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
+ or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
+ list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
+ latents as `image`, but if passing latents directly it is not encoded again.
+ prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+ instead.
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+ not greater than `1`).
+ true_cfg_scale (`float`, *optional*, defaults to 6.0):
+ When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The height in pixels of the generated image. This is set to 1024 by default for the best results.
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The width in pixels of the generated image. This is set to 1024 by default for the best results.
+ num_inference_steps (`int`, *optional*, defaults to 28):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+ will be used.
+ guidance_scale (`float`, *optional*, defaults to 3.5):
+ Guidance scale as defined in [Classifier-Free Diffusion
+ Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+ of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+ `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
+ the text `prompt`, usually at the expense of lower image quality.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+ to make generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+ tensor will be generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generate image. Choose between
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~pipelines.step1x_edit.Step1XEditPipelineOutput`] instead of a plain tuple.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ callback_on_step_end (`Callable`, *optional*):
+ A function that calls at the end of each denoising steps during the inference. The function is called
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+ `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+ `._callback_tensor_inputs` attribute of your pipeline class.
+
+ Examples:
+
+ Returns:
+ [`~pipelines.step1x_edit.Step1XEditPipelineOutput`] or `tuple`:
+ [`~pipelines.step1x_edit.Step1XEditPipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
+ returning a tuple, the first element is a list with the generated images.
+ """
+
+ device = self._execution_device
+
+ # 1. Preprocess image
+ image, ref_image, img_info, width, height = self.encode_image(
+ image=image[0] if isinstance(image, list) and len(image) > 0 else None,
+ width=width,
+ height=height,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt
+ )
+
+ # 2. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt=negative_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._joint_attention_kwargs = joint_attention_kwargs
+ self._current_timestep = None
+ self._interrupt = False
+
+ # 3. Define call parameters
+ 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]
+
+ lora_scale = (
+ self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+ )
+ has_neg_prompt = negative_prompt is not None or (
+ negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+ )
+ if not has_neg_prompt:
+ negative_prompt = "" if image is not None else "worst quality, wrong limbs, unreasonable limbs, normal quality, low quality, low res, blurry, text, watermark, logo, banner, extra digits, cropped, jpeg artifacts, signature, username, error, sketch ,duplicate, ugly, monochrome, horror, geometry, mutation, disgusting"
+ do_true_cfg = true_cfg_scale > 1
+ negative_text_ids = None
+ if (
+ image is None
+ and do_true_cfg
+ and prompt_embeds is None
+ and negative_prompt_embeds is None
+ and isinstance(prompt, str)
+ and isinstance(negative_prompt, str)
+ ):
+ combined_prompts = [prompt, negative_prompt]
+ combined_ref_images = [ref_image, ref_image]
+ (
+ combined_prompt_embeds,
+ combined_prompt_embeds_mask,
+ text_ids,
+ ) = self.encode_prompt(
+ ref_image=combined_ref_images,
+ prompt=combined_prompts,
+ prompt_embeds=None,
+ prompt_embeds_mask=None,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+ prompt_embeds, negative_prompt_embeds = combined_prompt_embeds.chunk(2, dim=0)
+ prompt_embeds_mask, negative_prompt_embeds_mask = combined_prompt_embeds_mask.chunk(2, dim=0)
+ negative_text_ids = text_ids
+ else:
+ (
+ prompt_embeds,
+ prompt_embeds_mask,
+ text_ids
+ ) = self.encode_prompt(
+ ref_image=ref_image,
+ prompt=prompt,
+ prompt_embeds=prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+ negative_text_ids = text_ids
+ if do_true_cfg:
+ (
+ negative_prompt_embeds,
+ negative_prompt_embeds_mask,
+ negative_text_ids,
+ ) = self.encode_prompt(
+ ref_image=ref_image,
+ prompt=negative_prompt,
+ prompt_embeds=negative_prompt_embeds,
+ prompt_embeds_mask=negative_prompt_embeds_mask,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+
+ if do_true_cfg and negative_text_ids is None:
+ negative_text_ids = text_ids
+
+ # 4. Prepare latent variables
+ num_channels_latents = self.transformer.config.in_channels // 4
+ latents, image_latents, latent_ids, image_ids = self.prepare_latents(
+ image,
+ batch_size * num_images_per_prompt,
+ num_channels_latents,
+ height,
+ width,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ latents,
+ )
+ if image_ids is not None:
+ latent_ids = torch.cat([latent_ids, image_ids], dim=0) # dim 0 is sequence dimension
+
+ # 5. Prepare timesteps
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+ image_seq_len = latents.shape[1]
+ mu = calculate_shift(
+ image_seq_len,
+ self.scheduler.config.get("base_image_seq_len", 256),
+ self.scheduler.config.get("max_image_seq_len", 4096),
+ self.scheduler.config.get("base_shift", 0.5),
+ self.scheduler.config.get("max_shift", 1.15),
+ )
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ self._num_timesteps = len(timesteps)
+
+ if self.transformer.config.guidance_embeds:
+ guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+ guidance = guidance.expand(latents.shape[0])
+ else:
+ guidance = None
+
+ if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+ negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+ ):
+ negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+ negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+ elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+ negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+ ):
+ ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+ ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+ if self.joint_attention_kwargs is None:
+ self._joint_attention_kwargs = {}
+
+ image_embeds = None
+ negative_image_embeds = None
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+ image_embeds = self.prepare_ip_adapter_image_embeds(
+ ip_adapter_image,
+ ip_adapter_image_embeds,
+ device,
+ batch_size * num_images_per_prompt,
+ )
+ if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+ negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+ negative_ip_adapter_image,
+ negative_ip_adapter_image_embeds,
+ device,
+ batch_size * num_images_per_prompt,
+ )
+
+ # 6. Denoising loop
+ # We set the index here to remove DtoH sync, helpful especially during compilation.
+ # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+ is_t2i = image_latents is None
+ self.scheduler.set_begin_index(0)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ self._current_timestep = t
+ if image_embeds is not None:
+ self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+
+ latent_model_input = latents
+ if image_latents is not None:
+ latent_model_input = torch.cat([latents, image_latents], dim=1)
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ if is_t2i and do_true_cfg:
+ # Reference implementation uses a dedicated t2i denoise path
+ # where cond/uncond are evaluated in one forward pass.
+ cfg_hidden_states = torch.cat([latent_model_input, latent_model_input], dim=0)
+ cfg_timestep = torch.cat([timestep, timestep], dim=0)
+ cfg_prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0)
+ cfg_prompt_embeds_mask = torch.cat([prompt_embeds_mask, negative_prompt_embeds_mask], dim=0)
+
+ cfg_guidance = guidance
+ if cfg_guidance is not None:
+ cfg_guidance = torch.cat([guidance, guidance], dim=0)
+
+ if image_embeds is not None and negative_image_embeds is not None:
+ self._joint_attention_kwargs["ip_adapter_image_embeds"] = [
+ torch.cat([img_embed, neg_img_embed], dim=0)
+ for img_embed, neg_img_embed in zip(image_embeds, negative_image_embeds)
+ ]
+
+ cfg_noise_pred = self.transformer(
+ hidden_states=cfg_hidden_states,
+ timestep=cfg_timestep / 1000,
+ guidance=cfg_guidance,
+ encoder_hidden_states=cfg_prompt_embeds,
+ prompt_embeds_mask=cfg_prompt_embeds_mask,
+ txt_ids=text_ids,
+ img_ids=latent_ids,
+ joint_attention_kwargs=self.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+ cfg_noise_pred = cfg_noise_pred[:, : latents.size(1)]
+ noise_pred, neg_noise_pred = cfg_noise_pred.chunk(2, dim=0)
+
+ if t.item() > timesteps_truncate:
+ diff = noise_pred - neg_noise_pred
+ diff_norm = torch.norm(diff, dim=(2), keepdim=True)
+ noise_pred = neg_noise_pred + true_cfg_scale * (
+ noise_pred - neg_noise_pred
+ ) / self.process_diff_norm(diff_norm, k=process_norm_power)
+ else:
+ noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+ else:
+ noise_pred = self.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ encoder_hidden_states=prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ txt_ids=text_ids,
+ img_ids=latent_ids,
+ joint_attention_kwargs=self.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+ noise_pred = noise_pred[:, : latents.size(1)]
+
+ if do_true_cfg:
+ if negative_image_embeds is not None:
+ self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
+ neg_noise_pred = self.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ encoder_hidden_states=negative_prompt_embeds,
+ prompt_embeds_mask=negative_prompt_embeds_mask,
+ txt_ids=negative_text_ids,
+ img_ids=latent_ids,
+ joint_attention_kwargs=self.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+ neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+ if t.item() > timesteps_truncate:
+ diff = noise_pred - neg_noise_pred
+ diff_norm = torch.norm(diff, dim=(2), keepdim=True)
+ noise_pred = neg_noise_pred + true_cfg_scale * (
+ noise_pred - neg_noise_pred
+ ) / self.process_diff_norm(diff_norm, k=process_norm_power)
+ else:
+ noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents_dtype = latents.dtype
+ if is_t2i:
+ t_prev = timesteps[i + 1] if i + 1 < len(timesteps) else latents.new_tensor(0.0)
+ # Timesteps are fed to the model in /1000 scale, so use the same
+ # normalized delta for manual t2i update.
+ dt = ((t_prev - t) / 1000).to(latents.dtype)
+ latents = latents + dt * noise_pred
+ else:
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+ if latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ latents = latents.to(latents_dtype)
+
+ 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, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ self._current_timestep = None
+
+ if output_type == "latent":
+ image = latents
+ else:
+ latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+ latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+ image = self.vae.decode(latents, return_dict=False)[0]
+ image = self.image_processor.postprocess(image, output_type=output_type)
+ image = self._output_process_image(image, img_info)
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (image,)
+
+ return Step1XEditPipelineOutput(images=image)
diff --git a/pipelines/step1x/transformer_step1x_edit.py b/pipelines/step1x/transformer_step1x_edit.py
new file mode 100644
index 000000000..f7130ef8c
--- /dev/null
+++ b/pipelines/step1x/transformer_step1x_edit.py
@@ -0,0 +1,1126 @@
+import inspect
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import math
+from functools import partial
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import (
+ Timesteps,
+ apply_rotary_emb,
+ get_1d_rotary_pos_embed,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _module_compute_dtype(module: nn.Module, fallback: torch.dtype) -> torch.dtype:
+ if hasattr(module, "sdnq_dequantizer") and hasattr(module.sdnq_dequantizer, "result_dtype"):
+ return module.sdnq_dequantizer.result_dtype
+ bias = getattr(module, "bias", None)
+ if isinstance(bias, torch.Tensor) and torch.is_floating_point(bias):
+ return bias.dtype
+ weight = getattr(module, "weight", None)
+ if isinstance(weight, torch.Tensor) and torch.is_floating_point(weight):
+ return weight.dtype
+ return fallback
+
+
+def _get_projections(attn: "Step1XEditAttention", hidden_states, encoder_hidden_states=None):
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ encoder_query = encoder_key = encoder_value = None
+ if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+ encoder_query = attn.add_q_proj(encoder_hidden_states)
+ encoder_key = attn.add_k_proj(encoder_hidden_states)
+ encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_fused_projections(attn: "Step1XEditAttention", hidden_states, encoder_hidden_states=None):
+ query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+
+ encoder_query = encoder_key = encoder_value = (None,)
+ if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
+ encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_qkv_projections(attn: "Step1XEditAttention", hidden_states, encoder_hidden_states=None):
+ if attn.fused_projections:
+ return _get_fused_projections(attn, hidden_states, encoder_hidden_states)
+ return _get_projections(attn, hidden_states, encoder_hidden_states)
+
+
+def apply_gate(x, gate=None, tanh=False):
+ """Applies a gating mechanism to the input tensor
+
+ Args:
+ x (torch.Tensor): input tensor.
+ gate (torch.Tensor, optional): gate tensor. Defaults to None.
+ tanh (bool, optional): whether to use tanh function. Defaults to False.
+
+ Returns:
+ torch.Tensor: the output tensor after apply gate.
+ """
+ if gate is None:
+ return x
+ if tanh:
+ return x * gate.unsqueeze(1).tanh()
+ else:
+ return x * gate.unsqueeze(1)
+
+
+class Step1XEditAttnProcessor:
+ _attention_backend = None
+
+ def __init__(self):
+
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
+
+ def __call__(
+ self,
+ attn: "Step1XEditAttention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
+ attn, hidden_states, encoder_hidden_states
+ )
+
+ query = query.unflatten(-1, (attn.heads, -1))
+ key = key.unflatten(-1, (attn.heads, -1))
+ value = value.unflatten(-1, (attn.heads, -1))
+
+ query = attn.norm_q(query)
+ key = attn.norm_k(key)
+
+ if attn.added_kv_proj_dim is not None:
+ encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
+ encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
+ encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+ encoder_query = attn.norm_added_q(encoder_query)
+ encoder_key = attn.norm_added_k(encoder_key)
+
+ query = torch.cat([encoder_query, query], dim=1)
+ key = torch.cat([encoder_key, key], dim=1)
+ value = torch.cat([encoder_value, value], dim=1)
+
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+ key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+ hidden_states = dispatch_attention_fn(
+ query, key, value, attn_mask=attention_mask, backend=self._attention_backend
+ )
+
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+ [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
+ )
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class Step1XEditAttention(torch.nn.Module, AttentionModuleMixin):
+ _default_processor_cls = Step1XEditAttnProcessor
+ _available_processors = [
+ Step1XEditAttnProcessor,
+ ]
+
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ out_bias: bool = True,
+ eps: float = 1e-6,
+ out_dim: Optional[int] = None,
+ context_pre_only: Optional[bool] = None,
+ pre_only: bool = False,
+ elementwise_affine: bool = True,
+ processor=None,
+ ):
+ super().__init__()
+
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.dropout = dropout
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.context_pre_only = context_pre_only
+ self.pre_only = pre_only
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.added_proj_bias = added_proj_bias
+
+ self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if not self.pre_only:
+ self.to_out = torch.nn.ModuleList([])
+ self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(torch.nn.Dropout(dropout))
+
+ if added_kv_proj_dim is not None:
+ self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
+
+ if processor is None:
+ processor = self._default_processor_cls()
+ self.set_processor(processor)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+ unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
+ if len(unused_kwargs) > 0:
+ logger.warning(
+ f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+ )
+ kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
+ return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
+
+
+@maybe_allow_in_graph
+class Step1XEditSingleTransformerBlock(nn.Module):
+ def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0):
+ super().__init__()
+ self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+ self.norm = AdaLayerNormZeroSingle(dim)
+ self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+ self.act_mlp = nn.GELU(approximate="tanh")
+ self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+ processor = Step1XEditAttnProcessor()
+
+ self.attn = Step1XEditAttention(
+ query_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=True,
+ processor=processor,
+ eps=1e-6,
+ pre_only=True,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ text_seq_len = encoder_hidden_states.shape[1]
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ residual = hidden_states
+ norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+ mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+ joint_attention_kwargs = joint_attention_kwargs or {}
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+ gate = gate.unsqueeze(1)
+ hidden_states = gate * self.proj_out(hidden_states)
+ hidden_states = residual + hidden_states
+ if hidden_states.dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+
+ return encoder_hidden_states, hidden_states
+
+
+@maybe_allow_in_graph
+class Step1XEditTransformerBlock(nn.Module):
+ def __init__(
+ self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+ ):
+ super().__init__()
+
+ self.norm1 = AdaLayerNormZero(dim)
+ self.norm1_context = AdaLayerNormZero(dim)
+
+ self.attn = Step1XEditAttention(
+ query_dim=dim,
+ added_kv_proj_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ context_pre_only=False,
+ bias=True,
+ processor=Step1XEditAttnProcessor(),
+ eps=eps,
+ )
+
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+ norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+ encoder_hidden_states, emb=temb
+ )
+ joint_attention_kwargs = joint_attention_kwargs or {}
+
+ # Attention.
+ attention_outputs = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_encoder_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ if len(attention_outputs) == 2:
+ attn_output, context_attn_output = attention_outputs
+ elif len(attention_outputs) == 3:
+ attn_output, context_attn_output, ip_attn_output = attention_outputs
+
+ # Process attention outputs for the `hidden_states`.
+ attn_output = gate_msa.unsqueeze(1) * attn_output
+ hidden_states = hidden_states + attn_output
+ norm_hidden_states = self.norm2(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+ ff_output = self.ff(norm_hidden_states)
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
+ hidden_states = hidden_states + ff_output
+
+ if len(attention_outputs) == 3:
+ hidden_states = hidden_states + ip_attn_output
+
+ # Process attention outputs for the `encoder_hidden_states`.
+ context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+ encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+ norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+ norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+ context_ff_output = self.ff_context(norm_encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+ if encoder_hidden_states.dtype == torch.float16:
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+ return encoder_hidden_states, hidden_states
+
+
+class Step1XEditPosEmbed(nn.Module):
+ # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+ def __init__(self, theta: int, axes_dim: List[int]):
+ super().__init__()
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ n_axes = ids.shape[-1]
+ cos_out = []
+ sin_out = []
+ pos = ids.float()
+ is_mps = ids.device.type == "mps"
+ is_npu = ids.device.type == "npu"
+ freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+ for i in range(n_axes):
+ cos, sin = get_1d_rotary_pos_embed(
+ self.axes_dim[i],
+ pos[:, i],
+ theta=self.theta,
+ repeat_interleave_real=True,
+ use_real=True,
+ freqs_dtype=freqs_dtype,
+ )
+ cos_out.append(cos)
+ sin_out.append(sin)
+ freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+ freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+ return freqs_cos, freqs_sin
+
+
+class Step1XEditMLP(nn.Module):
+ """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+ def __init__(
+ self,
+ in_channels,
+ hidden_channels=None,
+ out_features=None,
+ act_layer=nn.GELU,
+ norm_layer=None,
+ bias=True,
+ drop=0.0,
+ use_conv=False,
+ device=None,
+ dtype=None,
+ ):
+ super().__init__()
+ out_features = out_features or in_channels
+ hidden_channels = hidden_channels or in_channels
+ bias = (bias, bias)
+ drop_probs = (drop, drop)
+ linear_layer = nn.Linear
+
+ self.fc1 = linear_layer(
+ in_channels, hidden_channels, bias=bias[0], device=device, dtype=dtype
+ )
+ self.act = act_layer()
+ self.drop1 = nn.Dropout(drop_probs[0])
+ self.norm = (
+ norm_layer(hidden_channels, device=device, dtype=dtype)
+ if norm_layer is not None
+ else nn.Identity()
+ )
+ self.fc2 = linear_layer(
+ hidden_channels, out_features, bias=bias[1], device=device, dtype=dtype
+ )
+ self.drop2 = nn.Dropout(drop_probs[1])
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop1(x)
+ x = self.norm(x)
+ x = self.fc2(x)
+ x = self.drop2(x)
+ return x
+
+
+class Step1XEditMLPEmbedder(nn.Module):
+ def __init__(self, in_dim: int, hidden_dim: int):
+ super().__init__()
+ self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+ self.silu = nn.SiLU()
+ self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+ self.gradient_checkpointing = False
+
+ def set_gradient_checkpointing(self, enable: bool):
+ self.gradient_checkpointing = enable
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Step1XEditCrossAttnBlock(torch.nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ heads_num,
+ mlp_width_ratio: str = 4.0,
+ mlp_drop_rate: float = 0.0,
+ qk_norm: bool = False,
+ qkv_bias: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ super().__init__()
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+
+ self.norm1 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6
+ )
+ self.norm1_2 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6
+ )
+ self.self_attn_q = nn.Linear(
+ hidden_size, hidden_size, bias=qkv_bias
+ )
+ self.self_attn_kv = nn.Linear(
+ hidden_size, hidden_size*2, bias=qkv_bias
+ )
+ qk_norm_layer = nn.LayerNorm
+ self.self_attn_q_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_k_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_proj = nn.Linear(
+ hidden_size, hidden_size, bias=qkv_bias
+ )
+
+ self.norm2 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6
+ )
+ act_layer = nn.SiLU
+
+ self.adaLN_modulation = nn.Sequential(
+ act_layer(),
+ nn.Linear(hidden_size, 2 * hidden_size, bias=True),
+ )
+ # Zero-initialize the modulation
+ nn.init.zeros_(self.adaLN_modulation[1].weight)
+ nn.init.zeros_(self.adaLN_modulation[1].bias)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.Tensor, # timestep_aware_representations + context_aware_representations
+ attn_mask: torch.Tensor = None,
+ y: torch.Tensor=None,
+
+ ):
+ gate_msa, _gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
+
+ norm_x = self.norm1(x)
+ norm_y = self.norm1_2(y)
+ q = self.self_attn_q(norm_x)
+ q = q.view(q.size(0), q.size(1), self.heads_num, -1).permute(0, 2, 1, 3).contiguous()
+ kv = self.self_attn_kv(norm_y)
+ k, v = kv.view(kv.size(0), kv.size(1), 2, self.heads_num, -1).permute(2, 0, 3, 1, 4).contiguous().unbind(0)
+ # Apply QK-Norm if needed
+ q = self.self_attn_q_norm(q).to(v)
+ k = self.self_attn_k_norm(k).to(v)
+
+ # Self-Attention
+ attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask).transpose(1, 2)
+ attn = attn.reshape(attn.size(0), attn.size(1), -1)
+
+ x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
+
+ return x
+
+
+class Step1XEditIndividualTokenRefinerBlock(torch.nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ heads_num,
+ mlp_width_ratio: str = 4.0,
+ mlp_drop_rate: float = 0.0,
+ qk_norm: bool = False,
+ qkv_bias: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ super().__init__()
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+ self.norm1 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6
+ )
+ self.self_attn_qkv = nn.Linear(
+ hidden_size, hidden_size * 3, bias=qkv_bias
+ )
+ qk_norm_layer = nn.LayerNorm
+ self.self_attn_q_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_k_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_proj = nn.Linear(
+ hidden_size, hidden_size, bias=qkv_bias
+ )
+ self.norm2 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6
+ )
+ act_layer = nn.SiLU
+ self.mlp = Step1XEditMLP(
+ in_channels=hidden_size,
+ hidden_channels=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=mlp_drop_rate,
+ )
+
+ self.adaLN_modulation = nn.Sequential(
+ act_layer(),
+ nn.Linear(hidden_size, 2 * hidden_size, bias=True),
+ )
+
+ # Zero-initialize the modulation
+ nn.init.zeros_(self.adaLN_modulation[1].weight)
+ nn.init.zeros_(self.adaLN_modulation[1].bias)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.Tensor, # timestep_aware_representations + context_aware_representations
+ attn_mask: torch.Tensor = None,
+ y: torch.Tensor = None,
+ ):
+ gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
+
+ norm_x = self.norm1(x)
+ qkv = self.self_attn_qkv(norm_x)
+ q, k, v = qkv.view(qkv.size(0), qkv.size(1), 3, self.heads_num, -1).permute(2, 0, 3, 1, 4).contiguous().unbind(0)
+ # Apply QK-Norm if needed
+ q = self.self_attn_q_norm(q).to(v)
+ k = self.self_attn_k_norm(k).to(v)
+
+ # Self-Attention
+ attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask).transpose(1, 2)
+ attn = attn.reshape(attn.size(0), attn.size(1), -1)
+
+ x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
+
+ # FFN Layer
+ x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
+
+ return x
+
+
+class Step1XEditIndividualTokenRefiner(torch.nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ heads_num,
+ depth,
+ mlp_width_ratio: float = 4.0,
+ mlp_drop_rate: float = 0.0,
+ qk_norm: bool = False,
+ qkv_bias: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ super().__init__()
+ self.blocks = nn.ModuleList(
+ [
+ Step1XEditIndividualTokenRefinerBlock(
+ hidden_size=hidden_size,
+ heads_num=heads_num,
+ mlp_width_ratio=mlp_width_ratio,
+ mlp_drop_rate=mlp_drop_rate,
+ qk_norm=qk_norm,
+ qkv_bias=qkv_bias,
+ )
+ for _ in range(depth)
+ ]
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.LongTensor,
+ mask: Optional[torch.Tensor] = None,
+ y:torch.Tensor=None,
+ ):
+ self_attn_mask = None
+ if mask is not None:
+ batch_size = mask.shape[0]
+ seq_len = mask.shape[1]
+ mask = mask.to(x.device)
+ # batch_size x 1 x seq_len x seq_len
+ self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat(
+ 1, 1, seq_len, 1
+ )
+ # batch_size x 1 x seq_len x seq_len
+ self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
+ # batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num
+ self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
+ # avoids self-attention weight being NaN for padding tokens
+ self_attn_mask[:, :, :, 0] = True
+
+ for block in self.blocks:
+ x = block(x, c, self_attn_mask,y)
+
+ return x
+
+
+class Step1XEditTimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+
+ def __init__(
+ self,
+ hidden_size,
+ act_layer,
+ frequency_embedding_size=256,
+ max_period=10000,
+ out_size=None,
+ dtype=None,
+ device=None,
+ ):
+ super().__init__()
+ self.frequency_embedding_size = frequency_embedding_size
+ self.max_period = max_period
+ if out_size is None:
+ out_size = hidden_size
+
+ self.mlp = nn.Sequential(
+ nn.Linear(
+ frequency_embedding_size, hidden_size, bias=True
+ ),
+ act_layer(),
+ nn.Linear(hidden_size, out_size, bias=True),
+ )
+ nn.init.normal_(self.mlp[0].weight, std=0.02) # type: ignore
+ nn.init.normal_(self.mlp[2].weight, std=0.02) # type: ignore
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+
+ Args:
+ t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ dim (int): the dimension of the output.
+ max_period (int): controls the minimum frequency of the embeddings.
+
+ Returns:
+ embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
+
+ .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ """
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(start=0, end=half, dtype=torch.float32)
+ / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat(
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+ )
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(
+ t, self.frequency_embedding_size, self.max_period
+ ).type(self.mlp[0].weight.dtype) # type: ignore
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class Step1XEditTextProjection(nn.Module):
+ """
+ Projects text embeddings. Also handles dropout for classifier-free guidance.
+
+ Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+ """
+
+ def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
+ super().__init__()
+ self.linear_1 = nn.Linear(
+ in_features=in_channels,
+ out_features=hidden_size,
+ bias=True,
+ )
+ self.act_1 = act_layer()
+ self.linear_2 = nn.Linear(
+ in_features=hidden_size,
+ out_features=hidden_size,
+ bias=True,
+ )
+
+ def forward(self, caption):
+ hidden_states = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+
+class Step1XEditSingleTokenRefiner(torch.nn.Module):
+ """
+ A single token refiner block for llm text embedding refine.
+ """
+ def __init__(
+ self,
+ in_channels,
+ hidden_size,
+ heads_num,
+ depth,
+ mlp_width_ratio: float = 4.0,
+ mlp_drop_rate: float = 0.0,
+ qk_norm: bool = False,
+ qkv_bias: bool = True,
+ attn_mode: str = "torch",
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ super().__init__()
+ self.attn_mode = attn_mode
+ assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
+
+ self.input_embedder = nn.Linear(
+ in_channels, hidden_size, bias=True
+ )
+
+ act_layer = nn.SiLU
+ # Build timestep embedding layer
+ self.t_embedder = Step1XEditTimestepEmbedder(hidden_size, act_layer)
+ # Build context embedding layer
+ self.c_embedder = Step1XEditTextProjection(
+ in_channels, hidden_size, act_layer
+ )
+
+ self.individual_token_refiner = Step1XEditIndividualTokenRefiner(
+ hidden_size=hidden_size,
+ heads_num=heads_num,
+ depth=depth,
+ mlp_width_ratio=mlp_width_ratio,
+ mlp_drop_rate=mlp_drop_rate,
+ qk_norm=qk_norm,
+ qkv_bias=qkv_bias,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ t: torch.LongTensor,
+ mask: Optional[torch.LongTensor] = None,
+ ):
+ compute_dtype = _module_compute_dtype(self.input_embedder, x.dtype)
+ x = x.to(dtype=compute_dtype)
+ t = t.to(dtype=x.dtype)
+
+ timestep_aware_representations = self.t_embedder(t)
+
+ if mask is None:
+ context_aware_representations = x.mean(dim=1)
+ else:
+ mask = mask.to(device=x.device, dtype=torch.bool)
+ mask_float = mask.unsqueeze(-1).to(dtype=x.dtype) # [b, s1, 1]
+ context_aware_representations = (x * mask_float).sum(
+ dim=1
+ ) / mask_float.sum(dim=1).clamp_min(1.0)
+ context_dtype = _module_compute_dtype(self.c_embedder.linear_1, x.dtype)
+ context_aware_representations = context_aware_representations.to(dtype=context_dtype)
+ context_aware_representations = self.c_embedder(context_aware_representations)
+ c = timestep_aware_representations + context_aware_representations
+
+ x = self.input_embedder(x)
+ x = self.individual_token_refiner(x, c, mask)
+
+ return x
+
+
+class Step1XEditConnector(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels=3584,
+ hidden_size=4096,
+ heads_num=32,
+ depth=2,
+ ):
+ super().__init__()
+
+ self.S = Step1XEditSingleTokenRefiner(in_channels=in_channels,hidden_size=hidden_size,heads_num=heads_num,depth=depth)
+ self.global_proj_out=nn.Linear(in_channels, 768)
+
+ def forward(self, x, t, mask):
+ compute_dtype = _module_compute_dtype(self.global_proj_out, x.dtype)
+ x = x.to(dtype=compute_dtype)
+ t = t.to(dtype=x.dtype)
+ mask = mask.to(device=x.device, dtype=torch.bool)
+
+ t = t * 1000
+ mask_float = mask.unsqueeze(-1).to(dtype=x.dtype) # [b, s1, 1]
+
+ x_mean = (x * mask_float).sum(
+ dim=1
+ ) / mask_float.sum(dim=1).clamp_min(1.0)
+ x_mean = x_mean.to(dtype=compute_dtype)
+
+ global_out = self.global_proj_out(x_mean)
+ encoder_hidden_states = self.S(x,t,mask)
+ return encoder_hidden_states, global_out
+
+
+class Step1XEditTransformer2DModel(
+ ModelMixin,
+ ConfigMixin,
+ PeftAdapterMixin,
+ FromOriginalModelMixin,
+ CacheMixin,
+ AttentionMixin,
+):
+ """
+ The Transformer model introduced in Step1X-Edit.
+
+ Reference: https://arxiv.org/abs/2504.17761
+
+ Args:
+ patch_size (`int`, defaults to `1`):
+ Patch size to turn the input data into small patches.
+ in_channels (`int`, defaults to `64`):
+ The number of channels in the input.
+ out_channels (`int`, *optional*, defaults to `None`):
+ The number of channels in the output. If not specified, it defaults to `in_channels`.
+ num_layers (`int`, defaults to `19`):
+ The number of layers of dual stream DiT blocks to use.
+ num_single_layers (`int`, defaults to `38`):
+ The number of layers of single stream DiT blocks to use.
+ attention_head_dim (`int`, defaults to `128`):
+ The number of dimensions to use for each attention head.
+ num_attention_heads (`int`, defaults to `24`):
+ The number of attention heads to use.
+ joint_attention_dim (`int`, defaults to `4096`):
+ The number of dimensions to use for the joint attention (embedding/channel dimension of
+ `encoder_hidden_states`).
+ pooled_projection_dim (`int`, defaults to `768`):
+ The number of dimensions to use for the pooled projection.
+ guidance_embeds (`bool`, defaults to `False`):
+ Whether to use guidance embeddings for guidance-distilled variant of the model.
+ axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ The dimensions to use for the rotary positional embeddings.
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["Step1XEditTransformerBlock", "Step1XEditSingleTransformerBlock"]
+ _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+ _repeated_blocks = ["Step1XEditTransformerBlock", "Step1XEditSingleTransformerBlock"]
+
+ @register_to_config
+ def __init__(
+ self,
+ patch_size: int = 1,
+ in_channels: int = 64,
+ out_channels: Optional[int] = None,
+ num_layers: int = 19,
+ num_single_layers: int = 38,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 24,
+ joint_attention_dim: int = 4096,
+ timestep_in_dim: int = 256,
+ vector_in_dim: int = 768,
+ connector_in_channels=3584,
+ connector_hidden_size=4096,
+ connector_heads_num=32,
+ connector_depth=2,
+ guidance_embeds: bool = False,
+ axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ ):
+ super().__init__()
+ self.out_channels = out_channels or in_channels
+ self.inner_dim = num_attention_heads * attention_head_dim
+
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.pos_embed = Step1XEditPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+ self.time_embed = Step1XEditMLPEmbedder(timestep_in_dim, self.inner_dim)
+ self.vec_embed = Step1XEditMLPEmbedder(vector_in_dim, self.inner_dim)
+
+ self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+ self.connector = Step1XEditConnector(
+ connector_in_channels,
+ connector_hidden_size,
+ connector_heads_num,
+ connector_depth,
+ )
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ Step1XEditTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.single_transformer_blocks = nn.ModuleList(
+ [
+ Step1XEditSingleTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_single_layers)
+ ]
+ )
+
+ self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ prompt_embeds_mask: torch.Tensor = None,
+ timestep: torch.LongTensor = None,
+ img_ids: torch.Tensor = None,
+ txt_ids: torch.Tensor = None,
+ guidance: torch.Tensor = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ return_dict: bool = True,
+ controlnet_blocks_repeat: bool = False,
+ ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ """
+ The [`Step1XEditTransformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+ Input `hidden_states`.
+ encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+ pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+ from the embeddings of input conditions.
+ timestep ( `torch.LongTensor`):
+ Used to indicate denoising step.
+ block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+ A list of tensors that if specified are added to the residuals of transformer blocks.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+ `tuple` where the first element is the sample tensor.
+ """
+ if joint_attention_kwargs is not None:
+ joint_attention_kwargs = joint_attention_kwargs.copy()
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ scale_lora_layers(self, lora_scale)
+ else:
+ if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+ )
+
+ encoder_hidden_states, y = self.connector(
+ encoder_hidden_states, timestep, prompt_embeds_mask
+ )
+ hidden_states = self.x_embedder(hidden_states)
+
+ temb = self.time_embed(self.time_proj(timestep * 1000).to(timestep))
+ temb = temb + self.vec_embed(y)
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ ids = torch.cat((txt_ids, img_ids), dim=0)
+ image_rotary_emb = self.pos_embed(ids)
+
+ for _index_block, block in enumerate(self.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,
+ temb,
+ image_rotary_emb,
+ joint_attention_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ for _index_block, block in enumerate(self.single_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,
+ temb,
+ image_rotary_emb,
+ joint_attention_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+
+ return Transformer2DModelOutput(sample=output)
diff --git a/pipelines/ultraflux/autoencoder_kl.py b/pipelines/ultraflux/autoencoder_kl.py
new file mode 100644
index 000000000..5a47ba970
--- /dev/null
+++ b/pipelines/ultraflux/autoencoder_kl.py
@@ -0,0 +1,391 @@
+# Code borrow from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl.py
+from typing import Dict, Optional, Tuple, Union
+
+import os
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from diffusers.models.attention_processor import (
+ ADDED_KV_ATTENTION_PROCESSORS,
+ CROSS_ATTENTION_PROCESSORS,
+ Attention,
+ AttentionProcessor,
+ AttnAddedKVProcessor,
+ AttnProcessor,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from pipelines.ultraflux.vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
+
+
+class AutoencoderUltraFluxKL(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+ for all models (such as downloading or saving).
+
+ Parameters:
+ in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+ out_channels (int, *optional*, defaults to 3): Number of channels in the output.
+ down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ Tuple of downsample block types.
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ Tuple of upsample block types.
+ block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+ Tuple of block output channels.
+ act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+ latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+ sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+ scaling_factor (`float`, *optional*, defaults to 0.18215):
+ The component-wise standard deviation of the trained latent space computed using the first batch of the
+ training set. This is used to scale the latent space to have unit variance when training the diffusion
+ model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+ diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+ / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+ Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
+ force_upcast (`bool`, *optional*, default to `True`):
+ If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+ can be fine-tuned / trained to a lower range without loosing too much precision in which case
+ `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+ stride (int, *optional*, defaults to 1): stride for VAE.
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+
+ @register_to_config
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
+ up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
+ block_out_channels: Tuple[int] = (64,),
+ layers_per_block: int = 1,
+ act_fn: str = "silu",
+ latent_channels: int = 4,
+ norm_num_groups: int = 32,
+ sample_size: int = 32,
+ scaling_factor: float = 0.18215,
+ shift_factor: Optional[float] = None,
+ latents_mean: Optional[Tuple[float]] = None,
+ latents_std: Optional[Tuple[float]] = None,
+ force_upcast: float = True,
+ use_quant_conv: bool = True,
+ use_post_quant_conv: bool = True,
+ stride: int = 1,
+ ):
+ super().__init__()
+
+ # pass init params to Encoder
+ self.encoder = Encoder(
+ in_channels=in_channels,
+ out_channels=latent_channels,
+ down_block_types=down_block_types,
+ block_out_channels=block_out_channels,
+ layers_per_block=layers_per_block,
+ act_fn=act_fn,
+ norm_num_groups=norm_num_groups,
+ double_z=True,
+ stride=stride,
+ )
+
+ # pass init params to Decoder
+ self.decoder = Decoder(
+ in_channels=latent_channels,
+ out_channels=out_channels,
+ up_block_types=up_block_types,
+ block_out_channels=block_out_channels,
+ layers_per_block=layers_per_block,
+ norm_num_groups=norm_num_groups,
+ act_fn=act_fn,
+ stride=stride,
+ )
+
+ self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None
+ self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None
+
+ self.use_slicing = False
+ self.use_tiling = False
+
+ # only relevant if vae tiling is enabled
+ self.tile_sample_min_size = self.config.sample_size
+ sample_size = (
+ self.config.sample_size[0]
+ if isinstance(self.config.sample_size, (list, tuple))
+ else self.config.sample_size
+ )
+ self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1) ))
+ self.tile_overlap_factor = 0.25
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ if isinstance(module, (Encoder, Decoder)):
+ module.gradient_checkpointing = value
+
+ def enable_tiling(self, use_tiling: bool = True):
+ r"""
+ Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+ compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+ processing larger images.
+ """
+ self.use_tiling = use_tiling
+
+ def disable_tiling(self):
+ r"""
+ Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+ decoding in one step.
+ """
+ self.enable_tiling(False)
+
+ def enable_slicing(self):
+ r"""
+ Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+ compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+ """
+ self.use_slicing = True
+
+ def disable_slicing(self):
+ r"""
+ Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+ decoding in one step.
+ """
+ self.use_slicing = False
+
+ @property
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
+ r"""
+ Returns:
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
+ indexed by its weight name.
+ """
+ # set recursively
+ processors = {}
+
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+ if hasattr(module, "get_processor"):
+ processors[f"{name}.processor"] = module.get_processor()
+
+ for sub_name, child in module.named_children():
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+ return processors
+
+ for name, module in self.named_children():
+ fn_recursive_add_processors(name, module, processors)
+
+ return processors
+
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+ r"""
+ Sets the attention processor to use to compute attention.
+
+ Parameters:
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
+ for **all** `Attention` layers.
+
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+ processor. This is strongly recommended when setting trainable attention processors.
+
+ """
+ count = len(self.attn_processors.keys())
+
+ if isinstance(processor, dict) and len(processor) != count:
+ raise ValueError(
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+ )
+
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+ if hasattr(module, "set_processor"):
+ if not isinstance(processor, dict):
+ module.set_processor(processor)
+ else:
+ module.set_processor(processor.pop(f"{name}.processor"))
+
+ for sub_name, child in module.named_children():
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+ for name, module in self.named_children():
+ fn_recursive_attn_processor(name, module, processor)
+
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+ def set_default_attn_processor(self):
+ """
+ Disables custom attention processors and sets the default attention implementation.
+ """
+ if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+ processor = AttnAddedKVProcessor()
+ elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+ processor = AttnProcessor()
+ else:
+ raise ValueError(
+ f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+ )
+
+ self.set_attn_processor(processor)
+
+ @apply_forward_hook
+ def encode(
+ self, x: torch.Tensor, return_dict: bool = True
+ ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ """
+ Encode a batch of images into latents.
+
+ Args:
+ x (`torch.Tensor`): Input batch of images.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+ Returns:
+ The latent representations of the encoded images. If `return_dict` is True, a
+ [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+ """
+ if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+ return self.tiled_encode(x, return_dict=return_dict)
+
+ if self.use_slicing and x.shape[0] > 1:
+ encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+ h = torch.cat(encoded_slices)
+ else:
+ h = self.encoder(x)
+
+ if self.quant_conv is not None:
+ moments = self.quant_conv(h)
+ else:
+ moments = h
+
+ posterior = DiagonalGaussianDistribution(moments)
+
+ if not return_dict:
+ return (posterior,)
+
+ return AutoencoderKLOutput(latent_dist=posterior)
+
+ def _decode(self, z: torch.Tensor, return_dict: bool = True, partitioned: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
+ return self.tiled_decode(z, return_dict=return_dict, partitioned=partitioned)
+
+ if self.post_quant_conv is not None:
+ z = self.post_quant_conv(z)
+
+ dec = self.decoder(z, partitioned=partitioned)
+
+ if not return_dict:
+ return (dec,)
+
+ return DecoderOutput(sample=dec)
+
+ @apply_forward_hook
+ def decode(
+ self, z: torch.FloatTensor, return_dict: bool = True, generator=None, partitioned: bool = True
+ ) -> Union[DecoderOutput, torch.FloatTensor]:
+ """
+ Decode a batch of images.
+
+ Args:
+ z (`torch.Tensor`): Input batch of latent vectors.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.vae.DecoderOutput`] or `tuple`:
+ If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+ returned.
+
+ """
+ if self.use_slicing and z.shape[0] > 1:
+ decoded_slices = [self._decode(z_slice, partitioned=partitioned).sample for z_slice in z.split(1)]
+ decoded = torch.cat(decoded_slices)
+ else:
+ decoded = self._decode(z, partitioned=partitioned).sample
+
+ if not return_dict:
+ return (decoded,)
+
+ return DecoderOutput(sample=decoded)
+
+ def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+ blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+ for y in range(blend_extent):
+ b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
+ return b
+
+ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+ blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+ for x in range(blend_extent):
+ b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
+ return b
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ sample_posterior: bool = False,
+ return_dict: bool = True,
+ generator: Optional[torch.Generator] = None,
+ partitioned: bool = True,
+ ) -> Union[DecoderOutput, torch.Tensor]:
+ r"""
+ Args:
+ sample (`torch.Tensor`): Input sample.
+ sample_posterior (`bool`, *optional*, defaults to `False`):
+ Whether to sample from the posterior.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+ """
+ x = sample
+ posterior = self.encode(x).latent_dist
+ if sample_posterior:
+ z = posterior.sample(generator=generator)
+ else:
+ z = posterior.mode()
+ dec = self.decode(z, partitioned=partitioned).sample
+
+ if not return_dict:
+ return (dec,)
+
+ return DecoderOutput(sample=dec)
+
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+ def fuse_qkv_projections(self):
+ """
+ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+ are fused. For cross-attention modules, key and value projection matrices are fused.
+
+
+
+ This API is 🧪 experimental.
+
+
+ """
+ self.original_attn_processors = None
+
+ for _, attn_processor in self.attn_processors.items():
+ if "Added" in str(attn_processor.__class__.__name__):
+ raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+ self.original_attn_processors = self.attn_processors
+
+ for module in self.modules():
+ if isinstance(module, Attention):
+ module.fuse_projections(fuse=True)
+
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+ def unfuse_qkv_projections(self):
+ """Disables the fused QKV projection if enabled.
+
+
+
+ This API is 🧪 experimental.
+
+
+
+ """
+ if self.original_attn_processors is not None:
+ self.set_attn_processor(self.original_attn_processors)
diff --git a/pipelines/ultraflux/pipeline_flux.py b/pipelines/ultraflux/pipeline_flux.py
new file mode 100644
index 000000000..ff630f136
--- /dev/null
+++ b/pipelines/ultraflux/pipeline_flux.py
@@ -0,0 +1,750 @@
+# Code borrow from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux.py
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FluxLoraLoaderMixin, FromSingleFileMixin
+from pipelines.ultraflux.autoencoder_kl import AutoencoderUltraFluxKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+ USE_PEFT_BACKEND,
+ is_torch_xla_available,
+ logging,
+ replace_example_docstring,
+ scale_lora_layers,
+ unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+
+
+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 FluxPipeline
+
+ >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+ >>> pipe.to("cuda")
+ >>> prompt = "A cat holding a sign that says hello world"
+ >>> # Depending on the variant being used, the pipeline call will slightly vary.
+ >>> # Refer to the pipeline documentation for more details.
+ >>> image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
+ >>> image.save("flux.png")
+ ```
+"""
+
+
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.16,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class UltraFluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
+ r"""
+ The Flux pipeline for text-to-image generation.
+
+ Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+ Args:
+ transformer ([`FluxTransformer2DModel`]):
+ Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+ scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+ vae ([`AutoencoderUltraFluxKL`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+ text_encoder ([`CLIPTextModel`]):
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+ text_encoder_2 ([`T5EncoderModel`]):
+ [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+ the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+ tokenizer (`CLIPTokenizer`):
+ Tokenizer of class
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+ tokenizer_2 (`T5TokenizerFast`):
+ Second Tokenizer of class
+ [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+ """
+
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+ _optional_components = []
+ _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+ def __init__(
+ self,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ vae: AutoencoderUltraFluxKL,
+ text_encoder: CLIPTextModel,
+ tokenizer: CLIPTokenizer,
+ text_encoder_2: T5EncoderModel,
+ tokenizer_2: T5TokenizerFast,
+ transformer: FluxTransformer2DModel,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ text_encoder_2=text_encoder_2,
+ tokenizer=tokenizer,
+ tokenizer_2=tokenizer_2,
+ transformer=transformer,
+ scheduler=scheduler,
+ )
+
+ self.vae_scale_factor = 32
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+ self.tokenizer_max_length = (
+ self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+ )
+ self.default_sample_size = 64
+
+
+ def _get_t5_prompt_embeds(
+ self,
+ prompt: Union[str, List[str], None] = None,
+ num_images_per_prompt: int = 1,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ device = device or self._execution_device
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ text_inputs = self.tokenizer_2(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_length=False,
+ return_overflowing_tokens=False,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+
+ dtype = self.text_encoder_2.dtype
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ _, seq_len, _ = prompt_embeds.shape
+
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+ 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)
+
+ return prompt_embeds
+
+ def _get_clip_prompt_embeds(
+ self,
+ prompt: Union[str, List[str]],
+ num_images_per_prompt: int = 1,
+ device: Optional[torch.device] = None,
+ ):
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=self.tokenizer_max_length,
+ truncation=True,
+ return_overflowing_tokens=False,
+ return_length=False,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {self.tokenizer_max_length} tokens: {removed_text}"
+ )
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+ # Use pooled output of CLIPTextModel
+ prompt_embeds = prompt_embeds.pooler_output
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+ return prompt_embeds
+
+ def encode_prompt(
+ self,
+ prompt: Union[str, List[str]],
+ prompt_2: Union[str, List[str]],
+ device: Optional[torch.device] = None,
+ num_images_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ max_sequence_length: int = 512,
+ lora_scale: Optional[float] = None,
+ ):
+ r"""
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+ used in all text-encoders
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
+ lora_scale (`float`, *optional*):
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+ """
+ device = device or self._execution_device
+
+ # set lora scale so that monkey patched LoRA
+ # function of text encoder can correctly access it
+ if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+ self._lora_scale = lora_scale
+
+ # dynamically adjust the LoRA scale
+ if self.text_encoder is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(self.text_encoder, lora_scale)
+ if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(self.text_encoder_2, lora_scale)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_2 = prompt_2 or prompt
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+ # We only use the pooled prompt output from the CLIPTextModel
+ pooled_prompt_embeds = self._get_clip_prompt_embeds(
+ prompt=prompt,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+ prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=prompt_2,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ if self.text_encoder is not None:
+ if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(self.text_encoder, lora_scale)
+
+ if self.text_encoder_2 is not None:
+ if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+ dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+ text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+ text_ids = text_ids.repeat(num_images_per_prompt, 1, 1)
+
+ return prompt_embeds, pooled_prompt_embeds, text_ids
+
+ def check_inputs(
+ self,
+ prompt,
+ prompt_2,
+ height,
+ width,
+ prompt_embeds=None,
+ pooled_prompt_embeds=None,
+ callback_on_step_end_tensor_inputs=None,
+ max_sequence_length=None,
+ ):
+ if height % 8 != 0 or width % 8 != 0:
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+ 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 {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt_2 is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+ elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+ raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+ if prompt_embeds is not None and pooled_prompt_embeds is None:
+ raise ValueError(
+ "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+ )
+
+ if max_sequence_length is not None and max_sequence_length > 512:
+ raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+ @staticmethod
+ def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+ latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+ latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+ latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+ latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+ latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
+ latent_image_ids = latent_image_ids.reshape(
+ batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
+ )
+
+ return latent_image_ids.to(device=device, dtype=dtype)
+
+ @staticmethod
+ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+ latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+ latents = latents.permute(0, 2, 4, 1, 3, 5)
+ latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+ return latents
+
+ @staticmethod
+ def _unpack_latents(latents, height, width, vae_scale_factor):
+ batch_size, _num_patches, channels = latents.shape
+
+ height = height // vae_scale_factor
+ width = width // vae_scale_factor
+
+ latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+ latents = latents.reshape(batch_size, channels // (2 * 2), height * 2, width * 2)
+
+ return latents
+
+ def prepare_latents(
+ self,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ height = 2 * (int(height) // self.vae_scale_factor)
+ width = 2 * (int(width) // self.vae_scale_factor)
+
+ shape = (batch_size, num_channels_latents, height, width)
+
+ if latents is not None:
+ latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
+ return latents.to(device=device, dtype=dtype), latent_image_ids
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+ latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
+
+ return latents, latent_image_ids
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def joint_attention_kwargs(self):
+ return self._joint_attention_kwargs
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Union[str, List[str], None] = None,
+ prompt_2: Optional[Union[str, List[str]]] = None,
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ num_inference_steps: int = 28,
+ timesteps: Optional[List[int]] = None,
+ guidance_scale: float = 7.0,
+ num_images_per_prompt: Optional[int] = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.FloatTensor] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ partitioned: bool = True,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ max_sequence_length: int = 512,
+ ):
+ r"""
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+ instead.
+ prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+ will be used instead
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The height in pixels of the generated image. This is set to 1024 by default for the best results.
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The width in pixels of the generated image. This is set to 1024 by default for the best results.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+ passed will be used. Must be in descending order.
+ guidance_scale (`float`, *optional*, defaults to 7.0):
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+ usually at the expense of lower image quality.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+ to make generation deterministic.
+ latents (`torch.FloatTensor`, *optional*):
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+ tensor will ge generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generate image. Choose between
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ callback_on_step_end (`Callable`, *optional*):
+ A function that calls at the end of each denoising steps during the inference. The function is called
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+ `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+ `._callback_tensor_inputs` attribute of your pipeline class.
+ max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+ Examples:
+
+ Returns:
+ [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+ is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+ images.
+ """
+
+ height = height or self.default_sample_size * self.vae_scale_factor
+ width = width or self.default_sample_size * self.vae_scale_factor
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ prompt_2,
+ height,
+ width,
+ prompt_embeds=prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ max_sequence_length=max_sequence_length,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._joint_attention_kwargs = joint_attention_kwargs
+ self._interrupt = False
+
+ # 2. Define call parameters
+ 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
+
+ lora_scale = (
+ self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+ )
+ (
+ prompt_embeds,
+ pooled_prompt_embeds,
+ text_ids,
+ ) = self.encode_prompt(
+ prompt=prompt,
+ prompt_2=prompt_2,
+ prompt_embeds=prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ lora_scale=lora_scale,
+ )
+
+ # 4. Prepare latent variables
+ num_channels_latents = self.transformer.config.in_channels // 4
+ latents, latent_image_ids = self.prepare_latents(
+ batch_size * num_images_per_prompt,
+ num_channels_latents,
+ height,
+ width,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ latents,
+ )
+
+ # 5. Prepare timesteps
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+ image_seq_len = latents.shape[1]
+ mu = calculate_shift(
+ image_seq_len,
+ self.scheduler.config.base_image_seq_len,
+ self.scheduler.config.max_image_seq_len,
+ self.scheduler.config.base_shift,
+ self.scheduler.config.max_shift,
+ )
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ device,
+ timesteps,
+ sigmas,
+ mu=mu,
+ )
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ self._num_timesteps = len(timesteps)
+
+ # 6. Denoising loop
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ # handle guidance
+ if self.transformer.config.guidance_embeds:
+ guidance = torch.tensor([guidance_scale], device=device)
+ guidance = guidance.expand(latents.shape[0])
+ else:
+ guidance = None
+
+ noise_pred = self.transformer(
+ hidden_states=latents,
+ # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+ timestep=timestep / 1000,
+ guidance=guidance,
+ pooled_projections=pooled_prompt_embeds,
+ encoder_hidden_states=prompt_embeds,
+ txt_ids=text_ids,
+ img_ids=latent_image_ids,
+ joint_attention_kwargs=self.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents_dtype = latents.dtype
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+ if latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ latents = latents.to(latents_dtype)
+
+ 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, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ if output_type == "latent":
+ image = latents
+
+ else:
+ latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+ latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+ # Support both custom AutoencoderUltraFluxKL (with `partitioned`) and diffusers' AutoencoderKL
+ decode_kwargs = {"return_dict": False}
+ try:
+ sig = inspect.signature(self.vae.decode)
+ if "partitioned" in sig.parameters:
+ decode_kwargs["partitioned"] = partitioned
+ elif partitioned:
+ logger.warning("partitioned=True requested but VAE.decode has no 'partitioned' arg; using standard decode.")
+ except (ValueError, TypeError):
+ # If signature inspection fails, fall back to safest call without partitioned
+ pass
+ decoded = self.vae.decode(latents, **decode_kwargs)
+ image = decoded[0] if isinstance(decoded, (list, tuple)) else decoded
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (image,)
+
+ return FluxPipelineOutput(images=image)
diff --git a/pipelines/ultraflux/transformer_flux.py b/pipelines/ultraflux/transformer_flux.py
new file mode 100644
index 000000000..3bc3ef414
--- /dev/null
+++ b/pipelines/ultraflux/transformer_flux.py
@@ -0,0 +1,986 @@
+# Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX 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.
+
+import inspect
+from typing import Any, Dict, List, Optional, Tuple, Union
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import (
+ CombinedTimestepGuidanceTextProjEmbeddings,
+ CombinedTimestepTextProjEmbeddings,
+ apply_rotary_emb)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _get_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ encoder_query = encoder_key = encoder_value = None
+ if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+ encoder_query = attn.add_q_proj(encoder_hidden_states)
+ encoder_key = attn.add_k_proj(encoder_hidden_states)
+ encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_fused_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
+ query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+
+ encoder_query = encoder_key = encoder_value = (None,)
+ if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
+ encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_qkv_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
+ if attn.fused_projections:
+ return _get_fused_projections(attn, hidden_states, encoder_hidden_states)
+ return _get_projections(attn, hidden_states, encoder_hidden_states)
+
+
+class FluxAttnProcessor:
+ _attention_backend = None
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
+
+ def __call__(
+ self,
+ attn: "FluxAttention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
+ attn, hidden_states, encoder_hidden_states
+ )
+
+ query = query.unflatten(-1, (attn.heads, -1))
+ key = key.unflatten(-1, (attn.heads, -1))
+ value = value.unflatten(-1, (attn.heads, -1))
+
+ query = attn.norm_q(query)
+ key = attn.norm_k(key)
+
+ if attn.added_kv_proj_dim is not None:
+ encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
+ encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
+ encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+
+ encoder_query = attn.norm_added_q(encoder_query)
+ encoder_key = attn.norm_added_k(encoder_key)
+
+ query = torch.cat([encoder_query, query], dim=1)
+ key = torch.cat([encoder_key, key], dim=1)
+ value = torch.cat([encoder_value, value], dim=1)
+
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+ key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+ hidden_states = dispatch_attention_fn(
+ query, key, value, attn_mask=attention_mask, backend=self._attention_backend
+ )
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+ [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
+ )
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FluxIPAdapterAttnProcessor(torch.nn.Module):
+ """Flux Attention processor for IP-Adapter."""
+
+ _attention_backend = None
+
+ def __init__(
+ self, hidden_size: int, cross_attention_dim: int, num_tokens=(4,), scale=1.0, device=None, dtype=None
+ ):
+ super().__init__()
+
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ if not isinstance(num_tokens, (tuple, list)):
+ num_tokens = [num_tokens]
+
+ if not isinstance(scale, list):
+ scale = [scale] * len(num_tokens)
+ if len(scale) != len(num_tokens):
+ raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.")
+ self.scale = scale
+
+ self.to_k_ip = nn.ModuleList(
+ [
+ nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype)
+ for _ in range(len(num_tokens))
+ ]
+ )
+ self.to_v_ip = nn.ModuleList(
+ [
+ nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype)
+ for _ in range(len(num_tokens))
+ ]
+ )
+
+ def __call__(
+ self,
+ attn: "FluxAttention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ip_hidden_states: Optional[List[torch.Tensor]] = None,
+ ip_adapter_masks: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size = hidden_states.shape[0]
+
+ query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
+ attn, hidden_states, encoder_hidden_states
+ )
+
+ query = query.unflatten(-1, (attn.heads, -1))
+ key = key.unflatten(-1, (attn.heads, -1))
+ value = value.unflatten(-1, (attn.heads, -1))
+
+ query = attn.norm_q(query)
+ key = attn.norm_k(key)
+ ip_query = query
+
+ if encoder_hidden_states is not None:
+ encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
+ encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
+ encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+
+ encoder_query = attn.norm_added_q(encoder_query)
+ encoder_key = attn.norm_added_k(encoder_key)
+
+ query = torch.cat([encoder_query, query], dim=1)
+ key = torch.cat([encoder_key, key], dim=1)
+ value = torch.cat([encoder_value, value], dim=1)
+
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+ key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+ hidden_states = dispatch_attention_fn(
+ query,
+ key,
+ value,
+ attn_mask=attention_mask,
+ dropout_p=0.0,
+ is_causal=False,
+ backend=self._attention_backend,
+ )
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+ [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
+ )
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ # IP-adapter
+ ip_attn_output = torch.zeros_like(hidden_states)
+
+ for current_ip_hidden_states, scale, to_k_ip, to_v_ip in zip(
+ ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip
+ ):
+ ip_key = to_k_ip(current_ip_hidden_states)
+ ip_value = to_v_ip(current_ip_hidden_states)
+
+ ip_key = ip_key.view(batch_size, -1, attn.heads, attn.head_dim)
+ ip_value = ip_value.view(batch_size, -1, attn.heads, attn.head_dim)
+
+ current_ip_hidden_states = dispatch_attention_fn(
+ ip_query,
+ ip_key,
+ ip_value,
+ attn_mask=None,
+ dropout_p=0.0,
+ is_causal=False,
+ backend=self._attention_backend,
+ )
+ current_ip_hidden_states = current_ip_hidden_states.reshape(batch_size, -1, attn.heads * attn.head_dim)
+ current_ip_hidden_states = current_ip_hidden_states.to(ip_query.dtype)
+ ip_attn_output += scale * current_ip_hidden_states
+
+ return hidden_states, encoder_hidden_states, ip_attn_output
+ else:
+ return hidden_states
+
+
+class FluxAttention(torch.nn.Module, AttentionModuleMixin):
+ _default_processor_cls = FluxAttnProcessor
+ _available_processors = [
+ FluxAttnProcessor,
+ FluxIPAdapterAttnProcessor,
+ ]
+
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ out_bias: bool = True,
+ eps: float = 1e-5,
+ out_dim: Optional[int] = None,
+ context_pre_only: Optional[bool] = None,
+ pre_only: bool = False,
+ elementwise_affine: bool = True,
+ processor=None,
+ ):
+ super().__init__()
+
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.dropout = dropout
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.context_pre_only = context_pre_only
+ self.pre_only = pre_only
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.added_proj_bias = added_proj_bias
+
+ self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if not self.pre_only:
+ self.to_out = torch.nn.ModuleList([])
+ self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(torch.nn.Dropout(dropout))
+
+ if added_kv_proj_dim is not None:
+ self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
+
+ if processor is None:
+ processor = self._default_processor_cls()
+ self.set_processor(processor)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+ unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
+ if len(unused_kwargs) > 0:
+ logger.warning(
+ f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+ )
+ kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
+ return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
+
+
+@maybe_allow_in_graph
+class FluxSingleTransformerBlock(nn.Module):
+ def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0):
+ super().__init__()
+ self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+ self.norm = AdaLayerNormZeroSingle(dim)
+ self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+ self.act_mlp = nn.GELU(approximate="tanh")
+ self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+ self.attn = FluxAttention(
+ query_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=True,
+ processor=FluxAttnProcessor(),
+ eps=1e-6,
+ pre_only=True,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ text_seq_len = encoder_hidden_states.shape[1]
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ residual = hidden_states
+ norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+ mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+ joint_attention_kwargs = joint_attention_kwargs or {}
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+ gate = gate.unsqueeze(1)
+ hidden_states = gate * self.proj_out(hidden_states)
+ hidden_states = residual + hidden_states
+ if hidden_states.dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+ return encoder_hidden_states, hidden_states
+
+
+@maybe_allow_in_graph
+class FluxTransformerBlock(nn.Module):
+ def __init__(
+ self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+ ):
+ super().__init__()
+
+ self.norm1 = AdaLayerNormZero(dim)
+ self.norm1_context = AdaLayerNormZero(dim)
+
+ self.attn = FluxAttention(
+ query_dim=dim,
+ added_kv_proj_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ context_pre_only=False,
+ bias=True,
+ processor=FluxAttnProcessor(),
+ eps=eps,
+ )
+
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+ norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+ encoder_hidden_states, emb=temb
+ )
+ joint_attention_kwargs = joint_attention_kwargs or {}
+
+ # Attention.
+ attention_outputs = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_encoder_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ if len(attention_outputs) == 2:
+ attn_output, context_attn_output = attention_outputs
+ elif len(attention_outputs) == 3:
+ attn_output, context_attn_output, ip_attn_output = attention_outputs
+
+ # Process attention outputs for the `hidden_states`.
+ attn_output = gate_msa.unsqueeze(1) * attn_output
+ hidden_states = hidden_states + attn_output
+
+ norm_hidden_states = self.norm2(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+ ff_output = self.ff(norm_hidden_states)
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+ hidden_states = hidden_states + ff_output
+ if len(attention_outputs) == 3:
+ hidden_states = hidden_states + ip_attn_output
+
+ # Process attention outputs for the `encoder_hidden_states`.
+ context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+ encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+ norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+ norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+ context_ff_output = self.ff_context(norm_encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+ if encoder_hidden_states.dtype == torch.float16:
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+ return encoder_hidden_states, hidden_states
+
+
+def find_correction_factor(num_rotations, dim, base, max_position_embeddings):
+ return (dim * math.log(max_position_embeddings/(num_rotations * 2 * math.pi)))/(2 * math.log(base)) #Inverse dim formula to find number of rotations
+
+
+def find_correction_range(low_ratio, high_ratio, dim, base, ori_max_pe_len):
+ """
+ Find the correction range for NTK-by-parts interpolation.
+ """
+ low = np.floor(find_correction_factor(low_ratio, dim, base, ori_max_pe_len))
+ high = np.ceil(find_correction_factor(high_ratio, dim, base, ori_max_pe_len))
+ return max(low, 0), min(high, dim-1) #Clamp values just in case
+
+
+def linear_ramp_mask(min, max, dim):
+ if min == max:
+ max += 0.001 #Prevent singularity
+
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+
+def find_newbase_ntk(dim, base, scale):
+ """
+ Calculate the new base for NTK-aware scaling.
+ """
+ return base * (scale ** (dim / (dim - 2)))
+
+
+def get_1d_rotary_pos_embed(
+ dim: int,
+ pos: Union[np.ndarray, int],
+ theta: float = 10000.0,
+ use_real=False,
+ linear_factor=1.0,
+ ntk_factor=1.0,
+ repeat_interleave_real=True,
+ freqs_dtype=torch.float32,
+ yarn=True,
+ max_pe_len=None,
+ ori_max_pe_len=64,
+ current_timestep=1.0,
+ resonance: bool = True,
+ resonance_min_rot: float = 1.0,
+):
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos)
+ device = pos.device
+ dtype = freqs_dtype
+
+ # Dimension indices (even positions) matching the original implementation.
+ idx = torch.arange(0, dim, 2, dtype=dtype, device=device) / dim # a_i = i / D
+
+ if yarn and max_pe_len is not None and max_pe_len > ori_max_pe_len:
+ if not isinstance(max_pe_len, torch.Tensor):
+ max_pe_len = torch.tensor(max_pe_len, dtype=dtype, device=device)
+ scale = torch.clamp_min(max_pe_len / ori_max_pe_len, 1.0) # s >= 1
+
+ # =======================
+ # === Resonance alignment ===
+ # =======================
+ if resonance:
+ L = torch.tensor(float(ori_max_pe_len), dtype=dtype, device=device)
+ theta_t = torch.tensor(theta, dtype=dtype, device=device)
+ inv_freq = torch.exp(-idx * torch.log(theta_t))
+ rot = (L * inv_freq) / (2.0 * math.pi) # r(a) = L * theta^{-a} / (2 * pi)
+ use_resonance = rot >= 1.0
+ rot_rounded = torch.round(rot)
+ k = torch.clamp(rot_rounded, min=resonance_min_rot)
+ target_inv_freq = (2.0 * math.pi * k) / L # = θ^{-alpha_res}
+ alpha_res = -torch.log(target_inv_freq) / torch.log(theta_t)
+ idx_eff = torch.where(use_resonance, alpha_res, idx)
+ else:
+ idx_eff = idx
+
+
+ # ---- YaRN basis spectra (using resonance-aligned idx_eff) ----
+ beta_0, beta_1 = 1.25, 0.75
+ gamma_0, gamma_1 = 16, 2
+
+ # base: 1 / theta^a
+ # freqs_base = 1.0 / torch.exp(idx_eff * torch.log(torch.tensor(theta, dtype=dtype, device=device))) # [D/2]
+ freqs_base = 1.0 / torch.exp(idx_eff * torch.log(torch.tensor(theta, dtype=dtype, device=device)))
+ # linear(PI): base / s
+ freqs_linear = freqs_base / scale
+
+ # NTK-aware: use new_base^a
+ new_base = find_newbase_ntk(dim, theta, scale)
+ if isinstance(new_base, torch.Tensor) and new_base.dim() > 0:
+ new_base = new_base.view(-1, 1)
+ new_base = torch.tensor(float(new_base), dtype=dtype, device=device)
+ freqs_ntk = 1.0 / torch.exp(idx_eff * torch.log(new_base)) # [D/2]
+
+ # ---- YaRN β band: linear ↔ NTK interpolation (same as original logic) ----
+ low, high = find_correction_range(beta_0, beta_1, dim, theta, ori_max_pe_len)
+ low = max(0, low)
+ high = min(dim // 2, high)
+ mask_beta = (1 - linear_ramp_mask(low, high, dim // 2).to(device).to(dtype))
+ freqs = freqs_linear * (1 - mask_beta) + freqs_ntk * mask_beta
+
+ # ---- YaRN gamma band: high frequencies fall back to the base spectrum (same as original logic) ----
+ low, high = find_correction_range(gamma_0, gamma_1, dim, theta, ori_max_pe_len)
+ low = max(0, low)
+ high = min(dim // 2, high)
+ mask_gamma = (1 - linear_ramp_mask(low, high, dim // 2).to(device).to(dtype))
+ freqs = freqs * (1 - mask_gamma) + freqs_base * mask_gamma
+
+ else:
+ # Within the training window or when YARN is disabled: keep the original RoPE (supports ntk_factor / linear_factor).
+ theta_ntk = theta * ntk_factor
+ idx0 = torch.arange(0, dim, 2, dtype=dtype, device=device) / dim
+ freqs = 1.0 / torch.exp(idx0 * torch.log(torch.tensor(theta_ntk, dtype=dtype, device=device)))
+ freqs = freqs / linear_factor
+
+ # Phase outer product.
+ freqs = torch.outer(pos if isinstance(pos, torch.Tensor) else torch.tensor(pos, device=device), freqs)
+
+ if freqs.device.type == "npu":
+ freqs = freqs.float()
+
+ if use_real and repeat_interleave_real:
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float()
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float()
+ if yarn and max_pe_len is not None and max_pe_len > ori_max_pe_len:
+ mscale = torch.where(scale <= 1., torch.tensor(1.0, device=scale.device, dtype=scale.dtype),
+ 0.1 * torch.log(scale) + 1.0)
+ freqs_cos = freqs_cos * mscale
+ freqs_sin = freqs_sin * mscale
+ return freqs_cos, freqs_sin
+ elif use_real:
+ return torch.cat([freqs.cos(), freqs.cos()], dim=-1).float(), torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()
+ else:
+ return torch.polar(torch.ones_like(freqs), freqs)
+
+
+
+class FluxPosEmbed(nn.Module):
+ def __init__(
+ self,
+ theta: int,
+ axes_dim: List[int],
+ method: str = 'yarn',
+ base_resolution_hw: Optional[Tuple[int, int]] = None, # Training (H, W) resolution in pixels used to decouple YARN.
+ axis_order: str = "t-h-w", # Axis semantics for ids; defaults to [time/text, H, W].
+ ):
+ super().__init__()
+ self.theta = theta
+ self.axes_dim = axes_dim
+ self.method = method
+
+ # Keep the original square-image defaults.
+ self.base_resolution = 2048
+ self.patch_size = 32
+ self.base_patches = self.base_resolution // self.patch_size # Fallback only.
+
+ # Vision-YARN keeps independent training baseline patch counts for H and W.
+ if base_resolution_hw is None:
+ # Preserve legacy behavior where H_train == W_train == base_resolution.
+ H_train_px, W_train_px = self.base_resolution, self.base_resolution
+ else:
+ H_train_px, W_train_px = base_resolution_hw
+
+ self.base_patches_hw = (
+ max(1, H_train_px // self.patch_size), # H_train / ps
+ max(1, W_train_px // self.patch_size), # W_train / ps
+ )
+
+ # Specify the axis semantics for ids (change this if your ordering differs).
+ axis_order = axis_order.lower().replace("_", "").replace("-", "")
+ if axis_order not in ("thw", "twh"):
+ raise ValueError("axis_order must be 't-h-w' or 't-w-h'")
+ self.axis_order = axis_order
+
+ @staticmethod
+ def _span_length(pos_1d: torch.Tensor) -> int:
+ # Robust span length = max - min + 1 to guard against offsets.
+ return int(pos_1d.max().item() - pos_1d.min().item() + 1)
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ """
+ ids: [S, n_axes], typically cat(txt_ids, img_ids) with n_axes == len(axes_dim).
+ For axis_order='t-h-w': i=0 -> text/time, i=1 -> H, i=2 -> W (use 't-w-h' if swapped).
+ """
+ n_axes = ids.shape[-1]
+ assert n_axes == len(self.axes_dim), "axes_dim must match the number of axes in ids"
+
+ cos_out, sin_out = [], []
+ pos = ids.float()
+ is_mps = ids.device.type == "mps"
+ is_npu = ids.device.type == "npu"
+ freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+
+ # Determine the reference patch count per axis based on axis_order (only affects image H/W axes).
+ # thw: i=1->H, i=2->W ; twh: i=1->W, i=2->H
+ if self.axis_order == "thw":
+ axis_to_base_len = {0: None, 1: self.base_patches_hw[0], 2: self.base_patches_hw[1]}
+ else: # "twh"
+ axis_to_base_len = {0: None, 1: self.base_patches_hw[1], 2: self.base_patches_hw[0]}
+
+ for i in range(n_axes):
+ common_kwargs = {
+ 'dim': self.axes_dim[i],
+ 'pos': pos[:, i],
+ 'theta': self.theta,
+ 'repeat_interleave_real': True,
+ 'use_real': True,
+ 'freqs_dtype': freqs_dtype,
+ }
+
+ # === Core: decide per axis whether to enable YARN using that axis's own training baseline length. ===
+ base_len = axis_to_base_len.get(i, None)
+
+ if self.method == 'yarn' and base_len is not None:
+ current_len = self._span_length(pos[:, i]) # Current patch count for this axis (H_cur or W_cur).
+ if current_len > base_len:
+ max_pe_len = torch.tensor(current_len, dtype=freqs_dtype, device=pos.device)
+ cos, sin = get_1d_rotary_pos_embed(
+ **common_kwargs,
+ yarn=True,
+ max_pe_len=max_pe_len,
+ ori_max_pe_len=base_len, # Each axis uses its own training baseline length (Vision-YARN key idea).
+ )
+ else:
+ # Still within the training window, so YARN is not required.
+ cos, sin = get_1d_rotary_pos_embed(**common_kwargs)
+ else:
+ # Non-image axes (such as t/text) or non-YARN methods stay unchanged.
+ cos, sin = get_1d_rotary_pos_embed(**common_kwargs)
+
+ cos_out.append(cos)
+ sin_out.append(sin)
+
+ freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+ freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+ return freqs_cos, freqs_sin
+
+
+class FluxTransformer2DModel(
+ ModelMixin,
+ ConfigMixin,
+ PeftAdapterMixin,
+ FromOriginalModelMixin,
+ FluxTransformer2DLoadersMixin,
+ CacheMixin,
+ AttentionMixin,
+):
+ """
+ The Transformer model introduced in Flux with Yarn support.
+
+ Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+ Args:
+ patch_size (`int`, defaults to `1`):
+ Patch size to turn the input data into small patches.
+ in_channels (`int`, defaults to `64`):
+ The number of channels in the input.
+ out_channels (`int`, *optional*, defaults to `None`):
+ The number of channels in the output. If not specified, it defaults to `in_channels`.
+ num_layers (`int`, defaults to `19`):
+ The number of layers of dual stream DiT blocks to use.
+ num_single_layers (`int`, defaults to `38`):
+ The number of layers of single stream DiT blocks to use.
+ attention_head_dim (`int`, defaults to `128`):
+ The number of dimensions to use for each attention head.
+ num_attention_heads (`int`, defaults to `24`):
+ The number of attention heads to use.
+ joint_attention_dim (`int`, defaults to `4096`):
+ The number of dimensions to use for the joint attention (embedding/channel dimension of
+ `encoder_hidden_states`).
+ pooled_projection_dim (`int`, defaults to `768`):
+ The number of dimensions to use for the pooled projection.
+ guidance_embeds (`bool`, defaults to `False`):
+ Whether to use guidance embeddings for guidance-distilled variant of the model.
+ axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ The dimensions to use for the rotary positional embeddings.
+ method (`str`, defaults to `'yarn'`):
+ Position encoding method. Options: 'base', 'ntk', 'yarn'
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
+ _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+ _repeated_blocks = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
+
+ @register_to_config
+ def __init__(
+ self,
+ patch_size: int = 1,
+ in_channels: int = 64,
+ out_channels: Optional[int] = None,
+ num_layers: int = 19,
+ num_single_layers: int = 38,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 24,
+ joint_attention_dim: int = 4096,
+ pooled_projection_dim: int = 768,
+ guidance_embeds: bool = False,
+ axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ method: str = 'yarn'
+ ):
+ super().__init__()
+ self.out_channels = out_channels or in_channels
+ self.inner_dim = num_attention_heads * attention_head_dim
+
+ self.pos_embed = FluxPosEmbed(
+ theta=10000,
+ axes_dim=axes_dims_rope,
+ method=method
+ )
+
+ text_time_guidance_cls = (
+ CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+ )
+ self.time_text_embed = text_time_guidance_cls(
+ embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+ )
+
+ self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ FluxTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.single_transformer_blocks = nn.ModuleList(
+ [
+ FluxSingleTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_single_layers)
+ ]
+ )
+
+ self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ pooled_projections: torch.Tensor = None,
+ timestep: torch.LongTensor = None,
+ img_ids: torch.Tensor = None,
+ txt_ids: torch.Tensor = None,
+ guidance: torch.Tensor = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ controlnet_block_samples=None,
+ controlnet_single_block_samples=None,
+ return_dict: bool = True,
+ controlnet_blocks_repeat: bool = False,
+ ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ """
+ The [`FluxTransformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+ Input `hidden_states`.
+ encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+ pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+ from the embeddings of input conditions.
+ timestep ( `torch.LongTensor`):
+ Used to indicate denoising step.
+ block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+ A list of tensors that if specified are added to the residuals of transformer blocks.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+ `tuple` where the first element is the sample tensor.
+ """
+ if joint_attention_kwargs is not None:
+ joint_attention_kwargs = joint_attention_kwargs.copy()
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ scale_lora_layers(self, lora_scale)
+ else:
+ if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+ )
+
+ hidden_states = self.x_embedder(hidden_states)
+ timestep = timestep.to(hidden_states.dtype) * 1000
+
+ if guidance is not None:
+ guidance = guidance.to(hidden_states.dtype) * 1000
+
+ temb = (
+ self.time_text_embed(timestep, pooled_projections)
+ if guidance is None
+ else self.time_text_embed(timestep, guidance, pooled_projections)
+ )
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ if txt_ids.ndim == 3:
+ logger.warning(
+ "Passing `txt_ids` 3d torch.Tensor is deprecated."
+ "Please remove the batch dimension and pass it as a 2d torch Tensor"
+ )
+ txt_ids = txt_ids[0]
+ if img_ids.ndim == 3:
+ logger.warning(
+ "Passing `img_ids` 3d torch.Tensor is deprecated."
+ "Please remove the batch dimension and pass it as a 2d torch Tensor"
+ )
+ img_ids = img_ids[0]
+
+ ids = torch.cat((txt_ids, img_ids), dim=0)
+ image_rotary_emb = self.pos_embed(ids)
+
+ if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+ ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+ ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+ joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+ for index_block, block in enumerate(self.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,
+ temb,
+ image_rotary_emb,
+ joint_attention_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ if controlnet_block_samples is not None:
+ interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ if controlnet_blocks_repeat:
+ hidden_states = (
+ hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+ )
+ else:
+ hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+ for index_block, block in enumerate(self.single_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,
+ temb,
+ image_rotary_emb,
+ joint_attention_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ if controlnet_single_block_samples is not None:
+ interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ hidden_states = hidden_states + controlnet_single_block_samples[index_block // interval_control]
+
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ if USE_PEFT_BACKEND:
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+
+ return Transformer2DModelOutput(sample=output)
diff --git a/pipelines/ultraflux/vae.py b/pipelines/ultraflux/vae.py
new file mode 100644
index 000000000..2cdf5a63d
--- /dev/null
+++ b/pipelines/ultraflux/vae.py
@@ -0,0 +1,1043 @@
+# Code borrow from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from diffusers.utils import BaseOutput, is_torch_version
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.models.unets.unet_2d_blocks import (
+ AutoencoderTinyBlock,
+ UNetMidBlock2D,
+ get_down_block,
+ get_up_block,
+)
+
+import torch.nn.functional as F
+
+
+@dataclass
+class EncoderOutput(BaseOutput):
+ r"""
+ Output of encoding method.
+
+ Args:
+ latent (`torch.Tensor` of shape `(batch_size, num_channels, latent_height, latent_width)`):
+ The encoded latent.
+ """
+
+ latent: torch.Tensor
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+ r"""
+ Output of decoding method.
+
+ Args:
+ sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`):
+ The decoded output sample from the last layer of the model.
+ """
+
+ sample: torch.Tensor
+ commit_loss: Optional[torch.FloatTensor] = None
+
+
+class Encoder(nn.Module):
+ r"""
+ The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available
+ options.
+ block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ The number of output channels for each block.
+ layers_per_block (`int`, *optional*, defaults to 2):
+ The number of layers per block.
+ norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups for normalization.
+ act_fn (`str`, *optional*, defaults to `"silu"`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ double_z (`bool`, *optional*, defaults to `True`):
+ Whether to double the number of output channels for the last block.
+ stride (int, *optional*, defaults to 1): stride for VAE.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ double_z: bool = True,
+ mid_block_add_attention=True,
+ stride=1,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = nn.Conv2d(
+ in_channels,
+ block_out_channels[0],
+ kernel_size=3,
+ stride=stride,
+ padding=1,
+ )
+
+ self.down_blocks = nn.ModuleList([])
+
+ # down
+ output_channel = block_out_channels[0]
+ for i, down_block_type in enumerate(down_block_types):
+ input_channel = output_channel
+ output_channel = block_out_channels[i]
+ is_final_block = i == len(block_out_channels) - 1
+
+ down_block = get_down_block(
+ down_block_type,
+ num_layers=self.layers_per_block,
+ in_channels=input_channel,
+ out_channels=output_channel,
+ add_downsample=not is_final_block,
+ resnet_eps=1e-6,
+ downsample_padding=0,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ attention_head_dim=output_channel,
+ temb_channels=None,
+ )
+ self.down_blocks.append(down_block)
+
+ # mid
+ self.mid_block = UNetMidBlock2D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default",
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=None,
+ add_attention=mid_block_add_attention,
+ )
+
+ # out
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+
+ conv_out_channels = 2 * out_channels if double_z else out_channels
+ self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1)
+
+ self.gradient_checkpointing = False
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ r"""The forward method of the `Encoder` class."""
+
+ sample = self.conv_in(sample)
+
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ # down
+ if is_torch_version(">=", "1.11.0"):
+ for down_block in self.down_blocks:
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(down_block), sample, use_reentrant=False
+ )
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block), sample, use_reentrant=False
+ )
+ else:
+ for down_block in self.down_blocks:
+ sample = torch.utils.checkpoint.checkpoint(create_custom_forward(down_block), sample)
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample)
+
+ else:
+ # down
+ for down_block in self.down_blocks:
+ sample = down_block(sample)
+
+ # middle
+ sample = self.mid_block(sample)
+
+ # post-process
+ sample = self.conv_norm_out(sample)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class Decoder(nn.Module):
+ r"""
+ The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
+ block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ The number of output channels for each block.
+ layers_per_block (`int`, *optional*, defaults to 2):
+ The number of layers per block.
+ norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups for normalization.
+ act_fn (`str`, *optional*, defaults to `"silu"`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ norm_type (`str`, *optional*, defaults to `"group"`):
+ The normalization type to use. Can be either `"group"` or `"spatial"`.
+ stride (int, *optional*, defaults to 1): stride for VAE.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ norm_type: str = "group", # group, spatial
+ mid_block_add_attention=True,
+ stride=1,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+ self.stride = stride
+
+ self.conv_in = nn.Conv2d(
+ in_channels,
+ block_out_channels[-1],
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.up_blocks = nn.ModuleList([])
+
+ temb_channels = in_channels if norm_type == "spatial" else None
+
+ # mid
+ self.mid_block = UNetMidBlock2D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=temb_channels,
+ add_attention=mid_block_add_attention,
+ )
+
+ # up
+ reversed_block_out_channels = list(reversed(block_out_channels))
+ output_channel = reversed_block_out_channels[0]
+ for i, up_block_type in enumerate(up_block_types):
+ prev_output_channel = output_channel
+ output_channel = reversed_block_out_channels[i]
+
+ is_final_block = i == len(block_out_channels) - 1
+
+ up_block = get_up_block(
+ up_block_type,
+ num_layers=self.layers_per_block + 1,
+ in_channels=prev_output_channel,
+ out_channels=output_channel,
+ prev_output_channel=None,
+ add_upsample=not is_final_block,
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ attention_head_dim=output_channel,
+ temb_channels=temb_channels,
+ resnet_time_scale_shift=norm_type,
+ )
+ self.up_blocks.append(up_block)
+ prev_output_channel = output_channel
+
+ # out
+ if norm_type == "spatial":
+ self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+ else:
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+ self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ latent_embeds: Optional[torch.Tensor] = None,
+ partitioned: bool = True,
+ ) -> torch.Tensor:
+ r"""The forward method of the `Decoder` class."""
+
+ sample = self.conv_in(sample)
+
+ upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ if is_torch_version(">=", "1.11.0"):
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(up_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ else:
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block), sample, latent_embeds
+ )
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
+ else:
+ # middle
+ sample = self.mid_block(sample, latent_embeds)
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = up_block(sample, latent_embeds)
+
+ # partitioned VAE F16
+ if self.stride > 1 and partitioned:
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+ sample = self.conv_act(sample)
+
+ overlap_size = 1 # because last conv kernel_size = 3
+ res = []
+ partitioned_height = sample.shape[2] // self.stride
+ partitioned_width = sample.shape[3] // self.stride
+
+ assert self.stride == 2 # only support stride = 2 for now
+ rows = []
+ for i in range(0, sample.shape[2], partitioned_height):
+ row = []
+ for j in range(0, sample.shape[3], partitioned_width):
+ partition = sample[:,:, max(i - overlap_size, 0) : min(i + partitioned_height + overlap_size, sample.shape[2]), max(j - overlap_size, 0) : min(j + partitioned_width + overlap_size, sample.shape[3])]
+
+ # Stride = 2 padding setup.
+ if i==0 and j==0:
+ partition = F.pad(partition, (1, 0, 1, 0), "constant", 0)
+ elif i==0 and j>0:
+ partition = F.pad(partition, (0, 1, 1, 0), "constant", 0)
+ elif i>0 and j==0:
+ partition = F.pad(partition, (1, 0, 0, 1), "constant", 0)
+ elif i>0 and j>0:
+ partition = F.pad(partition, (0, 1, 0, 1), "constant", 0)
+
+ partition = F.interpolate(partition, scale_factor=self.stride, mode='nearest')
+ partition = self.conv_out(partition)
+ partition = partition[:,:,overlap_size:partitioned_height*2+overlap_size,overlap_size:partitioned_width*2+overlap_size]
+
+ row.append(partition)
+ rows.append(row)
+
+ for row in rows:
+ res.append(torch.cat(row, dim=3))
+
+ sample = torch.cat(res, dim=2)
+
+ # F16 VAE without partition
+ elif self.stride > 1:
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+
+ # add upsample
+ sample = F.interpolate(sample, scale_factor=self.stride, mode='nearest')
+
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ # F8 VAE
+ else:
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class UpSample(nn.Module):
+ r"""
+ The `UpSample` layer of a variational autoencoder that upsamples its input.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ ) -> None:
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ r"""The forward method of the `UpSample` class."""
+ x = torch.relu(x)
+ x = self.deconv(x)
+ return x
+
+
+class MaskConditionEncoder(nn.Module):
+ """
+ used in AsymmetricAutoencoderKL
+ """
+
+ def __init__(
+ self,
+ in_ch: int,
+ out_ch: int = 192,
+ res_ch: int = 768,
+ stride: int = 16,
+ ) -> None:
+ super().__init__()
+
+ channels = []
+ while stride > 1:
+ stride = stride // 2
+ in_ch_ = out_ch * 2
+ if out_ch > res_ch:
+ out_ch = res_ch
+ if stride == 1:
+ in_ch_ = res_ch
+ channels.append((in_ch_, out_ch))
+ out_ch *= 2
+
+ out_channels = []
+ for _in_ch, _out_ch in channels:
+ out_channels.append(_out_ch)
+ out_channels.append(channels[-1][0])
+
+ layers = []
+ in_ch_ = in_ch
+ for l in range(len(out_channels)):
+ out_ch_ = out_channels[l]
+ if l == 0 or l == 1:
+ layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=3, stride=1, padding=1))
+ else:
+ layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=4, stride=2, padding=1))
+ in_ch_ = out_ch_
+
+ self.layers = nn.Sequential(*layers)
+
+ def forward(self, x: torch.Tensor, mask=None) -> torch.Tensor:
+ r"""The forward method of the `MaskConditionEncoder` class."""
+ out = {}
+ for l in range(len(self.layers)):
+ layer = self.layers[l]
+ x = layer(x)
+ out[str(tuple(x.shape))] = x
+ x = torch.relu(x)
+ return out
+
+
+class MaskConditionDecoder(nn.Module):
+ r"""The `MaskConditionDecoder` should be used in combination with [`AsymmetricAutoencoderKL`] to enhance the model's
+ decoder with a conditioner on the mask and masked image.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
+ block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ The number of output channels for each block.
+ layers_per_block (`int`, *optional*, defaults to 2):
+ The number of layers per block.
+ norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups for normalization.
+ act_fn (`str`, *optional*, defaults to `"silu"`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ norm_type (`str`, *optional*, defaults to `"group"`):
+ The normalization type to use. Can be either `"group"` or `"spatial"`.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ norm_type: str = "group", # group, spatial
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = nn.Conv2d(
+ in_channels,
+ block_out_channels[-1],
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.up_blocks = nn.ModuleList([])
+
+ temb_channels = in_channels if norm_type == "spatial" else None
+
+ # mid
+ self.mid_block = UNetMidBlock2D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=temb_channels,
+ )
+
+ # up
+ reversed_block_out_channels = list(reversed(block_out_channels))
+ output_channel = reversed_block_out_channels[0]
+ for i, up_block_type in enumerate(up_block_types):
+ prev_output_channel = output_channel
+ output_channel = reversed_block_out_channels[i]
+
+ is_final_block = i == len(block_out_channels) - 1
+
+ up_block = get_up_block(
+ up_block_type,
+ num_layers=self.layers_per_block + 1,
+ in_channels=prev_output_channel,
+ out_channels=output_channel,
+ prev_output_channel=None,
+ add_upsample=not is_final_block,
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ attention_head_dim=output_channel,
+ temb_channels=temb_channels,
+ resnet_time_scale_shift=norm_type,
+ )
+ self.up_blocks.append(up_block)
+ prev_output_channel = output_channel
+
+ # condition encoder
+ self.condition_encoder = MaskConditionEncoder(
+ in_ch=out_channels,
+ out_ch=block_out_channels[0],
+ res_ch=block_out_channels[-1],
+ )
+
+ # out
+ if norm_type == "spatial":
+ self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+ else:
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+ self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ z: torch.Tensor,
+ image: Optional[torch.Tensor] = None,
+ mask: Optional[torch.Tensor] = None,
+ latent_embeds: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ r"""The forward method of the `MaskConditionDecoder` class."""
+ sample = z
+ sample = self.conv_in(sample)
+
+ upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ if is_torch_version(">=", "1.11.0"):
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ sample = sample.to(upscale_dtype)
+
+ # condition encoder
+ if image is not None and mask is not None:
+ masked_image = (1 - mask) * image
+ im_x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.condition_encoder),
+ masked_image,
+ mask,
+ use_reentrant=False,
+ )
+
+ # up
+ for up_block in self.up_blocks:
+ if image is not None and mask is not None:
+ sample_ = im_x[str(tuple(sample.shape))]
+ mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest")
+ sample = sample * mask_ + sample_ * (1 - mask_)
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(up_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ if image is not None and mask is not None:
+ sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask)
+ else:
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block), sample, latent_embeds
+ )
+ sample = sample.to(upscale_dtype)
+
+ # condition encoder
+ if image is not None and mask is not None:
+ masked_image = (1 - mask) * image
+ im_x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.condition_encoder),
+ masked_image,
+ mask,
+ )
+
+ # up
+ for up_block in self.up_blocks:
+ if image is not None and mask is not None:
+ sample_ = im_x[str(tuple(sample.shape))]
+ mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest")
+ sample = sample * mask_ + sample_ * (1 - mask_)
+ sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
+ if image is not None and mask is not None:
+ sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask)
+ else:
+ # middle
+ sample = self.mid_block(sample, latent_embeds)
+ sample = sample.to(upscale_dtype)
+
+ # condition encoder
+ if image is not None and mask is not None:
+ masked_image = (1 - mask) * image
+ im_x = self.condition_encoder(masked_image, mask)
+
+ # up
+ for up_block in self.up_blocks:
+ if image is not None and mask is not None:
+ sample_ = im_x[str(tuple(sample.shape))]
+ mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest")
+ sample = sample * mask_ + sample_ * (1 - mask_)
+ sample = up_block(sample, latent_embeds)
+ if image is not None and mask is not None:
+ sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask)
+
+ # post-process
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class VectorQuantizer(nn.Module):
+ """
+ Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly avoids costly matrix
+ multiplications and allows for post-hoc remapping of indices.
+ """
+
+ # NOTE: due to a bug the beta term was applied to the wrong term. for
+ # backwards compatibility we use the buggy version by default, but you can
+ # specify legacy=False to fix it.
+ def __init__(
+ self,
+ n_e: int,
+ vq_embed_dim: int,
+ beta: float,
+ remap=None,
+ unknown_index: str = "random",
+ sane_index_shape: bool = False,
+ legacy: bool = True,
+ ):
+ super().__init__()
+ self.n_e = n_e
+ self.vq_embed_dim = vq_embed_dim
+ self.beta = beta
+ self.legacy = legacy
+
+ self.embedding = nn.Embedding(self.n_e, self.vq_embed_dim)
+ self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+
+ self.remap = remap
+ if self.remap is not None:
+ self.register_buffer("used", torch.tensor(np.load(self.remap)))
+ self.used: torch.Tensor
+ self.re_embed = self.used.shape[0]
+ self.unknown_index = unknown_index # "random" or "extra" or integer
+ if self.unknown_index == "extra":
+ self.unknown_index = self.re_embed
+ self.re_embed = self.re_embed + 1
+ else:
+ self.re_embed = n_e
+
+ self.sane_index_shape = sane_index_shape
+
+ def remap_to_used(self, inds: torch.LongTensor) -> torch.LongTensor:
+ ishape = inds.shape
+ assert len(ishape) > 1
+ inds = inds.reshape(ishape[0], -1)
+ used = self.used.to(inds)
+ match = (inds[:, :, None] == used[None, None, ...]).long()
+ new = match.argmax(-1)
+ unknown = match.sum(2) < 1
+ if self.unknown_index == "random":
+ new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device)
+ else:
+ new[unknown] = self.unknown_index
+ return new.reshape(ishape)
+
+ def unmap_to_all(self, inds: torch.LongTensor) -> torch.LongTensor:
+ ishape = inds.shape
+ assert len(ishape) > 1
+ inds = inds.reshape(ishape[0], -1)
+ used = self.used.to(inds)
+ if self.re_embed > self.used.shape[0]: # extra token
+ inds[inds >= self.used.shape[0]] = 0 # simply set to zero
+ back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds)
+ return back.reshape(ishape)
+
+ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
+ # reshape z -> (batch, height, width, channel) and flatten
+ z = z.permute(0, 2, 3, 1).contiguous()
+ z_flattened = z.view(-1, self.vq_embed_dim)
+
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+ min_encoding_indices = torch.argmin(torch.cdist(z_flattened, self.embedding.weight), dim=1)
+
+ z_q = self.embedding(min_encoding_indices).view(z.shape)
+ perplexity = None
+ min_encodings = None
+
+ # compute loss for embedding
+ if not self.legacy:
+ loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
+ else:
+ loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
+
+ # preserve gradients
+ z_q: torch.Tensor = z + (z_q - z).detach()
+
+ # reshape back to match original input shape
+ z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+ if self.remap is not None:
+ min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis
+ min_encoding_indices = self.remap_to_used(min_encoding_indices)
+ min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten
+
+ if self.sane_index_shape:
+ min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3])
+
+ return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
+
+ def get_codebook_entry(self, indices: torch.LongTensor, shape: Tuple[int, ...]) -> torch.Tensor:
+ # shape specifying (batch, height, width, channel)
+ if self.remap is not None:
+ indices = indices.reshape(shape[0], -1) # add batch axis
+ indices = self.unmap_to_all(indices)
+ indices = indices.reshape(-1) # flatten again
+
+ # get quantized latent vectors
+ z_q: torch.Tensor = self.embedding(indices)
+
+ if shape is not None:
+ z_q = z_q.view(shape)
+ # reshape back to match original input shape
+ z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+ return z_q
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.deterministic = deterministic
+ self.std = torch.exp(0.5 * self.logvar)
+ self.var = torch.exp(self.logvar)
+ if self.deterministic:
+ self.var = self.std = torch.zeros_like(
+ self.mean, device=self.parameters.device, dtype=self.parameters.dtype
+ )
+
+ def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ # make sure sample is on the same device as the parameters and has same dtype
+ sample = randn_tensor(
+ self.mean.shape,
+ generator=generator,
+ device=self.parameters.device,
+ dtype=self.parameters.dtype,
+ )
+ x = self.mean + self.std * sample
+ return x
+
+ def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+ if self.deterministic:
+ return torch.Tensor([0.0])
+ else:
+ if other is None:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+ dim=[1, 2, 3],
+ )
+ else:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean - other.mean, 2) / other.var
+ + self.var / other.var
+ - 1.0
+ - self.logvar
+ + other.logvar,
+ dim=[1, 2, 3],
+ )
+
+ def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+ if self.deterministic:
+ return torch.Tensor([0.0])
+ logtwopi = np.log(2.0 * np.pi)
+ return 0.5 * torch.sum(
+ logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+ dim=dims,
+ )
+
+ def mode(self) -> torch.Tensor:
+ return self.mean
+
+
+class EncoderTiny(nn.Module):
+ r"""
+ The `EncoderTiny` layer is a simpler version of the `Encoder` layer.
+
+ Args:
+ in_channels (`int`):
+ The number of input channels.
+ out_channels (`int`):
+ The number of output channels.
+ num_blocks (`Tuple[int, ...]`):
+ Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to
+ use.
+ block_out_channels (`Tuple[int, ...]`):
+ The number of output channels for each block.
+ act_fn (`str`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_blocks: Tuple[int, ...],
+ block_out_channels: Tuple[int, ...],
+ act_fn: str,
+ ):
+ super().__init__()
+
+ layers = []
+ for i, num_block in enumerate(num_blocks):
+ num_channels = block_out_channels[i]
+
+ if i == 0:
+ layers.append(nn.Conv2d(in_channels, num_channels, kernel_size=3, padding=1))
+ else:
+ layers.append(
+ nn.Conv2d(
+ num_channels,
+ num_channels,
+ kernel_size=3,
+ padding=1,
+ stride=2,
+ bias=False,
+ )
+ )
+
+ for _ in range(num_block):
+ layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn))
+
+ layers.append(nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, padding=1))
+
+ self.layers = nn.Sequential(*layers)
+ self.gradient_checkpointing = False
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ r"""The forward method of the `EncoderTiny` class."""
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ if is_torch_version(">=", "1.11.0"):
+ x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x, use_reentrant=False)
+ else:
+ x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x)
+
+ else:
+ # scale image from [-1, 1] to [0, 1] to match TAESD convention
+ x = self.layers(x.add(1).div(2))
+
+ return x
+
+
+class DecoderTiny(nn.Module):
+ r"""
+ The `DecoderTiny` layer is a simpler version of the `Decoder` layer.
+
+ Args:
+ in_channels (`int`):
+ The number of input channels.
+ out_channels (`int`):
+ The number of output channels.
+ num_blocks (`Tuple[int, ...]`):
+ Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to
+ use.
+ block_out_channels (`Tuple[int, ...]`):
+ The number of output channels for each block.
+ upsampling_scaling_factor (`int`):
+ The scaling factor to use for upsampling.
+ act_fn (`str`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_blocks: Tuple[int, ...],
+ block_out_channels: Tuple[int, ...],
+ upsampling_scaling_factor: int,
+ act_fn: str,
+ upsample_fn: str,
+ ):
+ super().__init__()
+
+ layers = [
+ nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=1),
+ get_activation(act_fn),
+ ]
+
+ for i, num_block in enumerate(num_blocks):
+ is_final_block = i == (len(num_blocks) - 1)
+ num_channels = block_out_channels[i]
+
+ for _ in range(num_block):
+ layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn))
+
+ if not is_final_block:
+ layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor, mode=upsample_fn))
+
+ conv_out_channel = num_channels if not is_final_block else out_channels
+ layers.append(
+ nn.Conv2d(
+ num_channels,
+ conv_out_channel,
+ kernel_size=3,
+ padding=1,
+ bias=is_final_block,
+ )
+ )
+
+ self.layers = nn.Sequential(*layers)
+ self.gradient_checkpointing = False
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ r"""The forward method of the `DecoderTiny` class."""
+ # Clamp.
+ x = torch.tanh(x / 3) * 3
+
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ if is_torch_version(">=", "1.11.0"):
+ x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x, use_reentrant=False)
+ else:
+ x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x)
+
+ else:
+ x = self.layers(x)
+
+ # scale image from [0, 1] to [-1, 1] to match diffusers convention
+ return x.mul(2).sub(1)
diff --git a/pipelines/vibe/__init__.py b/pipelines/vibe/__init__.py
new file mode 100644
index 000000000..49b0390cb
--- /dev/null
+++ b/pipelines/vibe/__init__.py
@@ -0,0 +1,6 @@
+"""VIBE pipeline components for SD.Next."""
+
+from .vibe_sana_editing import VIBESanaEditingModel
+from .vibe_sana_pipeline import VIBESanaEditingPipeline, VIBESanaImagePipeline
+
+__all__ = ["VIBESanaEditingModel", "VIBESanaEditingPipeline", "VIBESanaImagePipeline"]
diff --git a/pipelines/vibe/aspects_multiscale.py b/pipelines/vibe/aspects_multiscale.py
new file mode 100644
index 000000000..29056c143
--- /dev/null
+++ b/pipelines/vibe/aspects_multiscale.py
@@ -0,0 +1,621 @@
+
+ASPECT_RATIO_384 = {
+ "0.25": [192.0, 768.0],
+ "0.26": [192.0, 744.0],
+ "0.27": [192.0, 720.0],
+ "0.28": [192.0, 696.0],
+ "0.32": [216.0, 672.0],
+ "0.33": [216.0, 648.0],
+ "0.35": [216.0, 624.0],
+ "0.4": [240.0, 600.0],
+ "0.42": [240.0, 576.0],
+ "0.48": [264.0, 552.0],
+ "0.5": [264.0, 528.0],
+ "0.52": [264.0, 504.0],
+ "0.57": [288.0, 504.0],
+ "0.6": [288.0, 480.0],
+ "0.68": [312.0, 456.0],
+ "0.72": [312.0, 432.0],
+ "0.78": [336.0, 432.0],
+ "0.82": [336.0, 408.0],
+ "0.88": [360.0, 408.0],
+ "0.94": [360.0, 384.0],
+ "1.0": [384.0, 384.0],
+ "1.07": [384.0, 360.0],
+ "1.13": [408.0, 360.0],
+ "1.21": [408.0, 336.0],
+ "1.29": [432.0, 336.0],
+ "1.38": [432.0, 312.0],
+ "1.46": [456.0, 312.0],
+ "1.67": [480.0, 288.0],
+ "1.75": [504.0, 288.0],
+ "2.0": [528.0, 264.0],
+ "2.09": [552.0, 264.0],
+ "2.4": [576.0, 240.0],
+ "2.5": [600.0, 240.0],
+ "2.89": [624.0, 216.0],
+ "3.0": [648.0, 216.0],
+ "3.11": [672.0, 216.0],
+ "3.62": [696.0, 192.0],
+ "3.75": [720.0, 192.0],
+ "3.88": [744.0, 192.0],
+ "4.0": [768.0, 192.0],
+}
+
+ASPECT_RATIO_512 = {
+ "0.25": [256.0, 1024.0],
+ "0.26": [256.0, 992.0],
+ "0.27": [256.0, 960.0],
+ "0.28": [256.0, 928.0],
+ "0.32": [288.0, 896.0],
+ "0.33": [288.0, 864.0],
+ "0.35": [288.0, 832.0],
+ "0.4": [320.0, 800.0],
+ "0.42": [320.0, 768.0],
+ "0.48": [352.0, 736.0],
+ "0.5": [352.0, 704.0],
+ "0.52": [352.0, 672.0],
+ "0.57": [384.0, 672.0],
+ "0.6": [384.0, 640.0],
+ "0.68": [416.0, 608.0],
+ "0.72": [416.0, 576.0],
+ "0.78": [448.0, 576.0],
+ "0.82": [448.0, 544.0],
+ "0.88": [480.0, 544.0],
+ "0.94": [480.0, 512.0],
+ "1.0": [512.0, 512.0],
+ "1.07": [512.0, 480.0],
+ "1.13": [544.0, 480.0],
+ "1.21": [544.0, 448.0],
+ "1.29": [576.0, 448.0],
+ "1.38": [576.0, 416.0],
+ "1.46": [608.0, 416.0],
+ "1.67": [640.0, 384.0],
+ "1.75": [672.0, 384.0],
+ "2.0": [704.0, 352.0],
+ "2.09": [736.0, 352.0],
+ "2.4": [768.0, 320.0],
+ "2.5": [800.0, 320.0],
+ "2.89": [832.0, 288.0],
+ "3.0": [864.0, 288.0],
+ "3.11": [896.0, 288.0],
+ "3.62": [928.0, 256.0],
+ "3.75": [960.0, 256.0],
+ "3.88": [992.0, 256.0],
+ "4.0": [1024.0, 256.0],
+}
+
+ASPECT_RATIO_768 = {
+ "0.25": [384.0, 1536.0],
+ "0.26": [384.0, 1488.0],
+ "0.27": [384.0, 1440.0],
+ "0.28": [384.0, 1392.0],
+ "0.32": [432.0, 1344.0],
+ "0.33": [432.0, 1296.0],
+ "0.35": [432.0, 1248.0],
+ "0.4": [480.0, 1200.0],
+ "0.42": [480.0, 1152.0],
+ "0.48": [528.0, 1104.0],
+ "0.5": [528.0, 1056.0],
+ "0.52": [528.0, 1008.0],
+ "0.57": [576.0, 1008.0],
+ "0.6": [576.0, 960.0],
+ "0.68": [624.0, 912.0],
+ "0.72": [624.0, 864.0],
+ "0.78": [672.0, 864.0],
+ "0.82": [672.0, 816.0],
+ "0.88": [720.0, 816.0],
+ "0.94": [720.0, 768.0],
+ "1.0": [768.0, 768.0],
+ "1.07": [768.0, 720.0],
+ "1.13": [816.0, 720.0],
+ "1.21": [816.0, 672.0],
+ "1.29": [864.0, 672.0],
+ "1.38": [864.0, 624.0],
+ "1.46": [912.0, 624.0],
+ "1.67": [960.0, 576.0],
+ "1.75": [1008.0, 576.0],
+ "2.0": [1056.0, 528.0],
+ "2.09": [1104.0, 528.0],
+ "2.4": [1152.0, 480.0],
+ "2.5": [1200.0, 480.0],
+ "2.89": [1248.0, 432.0],
+ "3.0": [1296.0, 432.0],
+ "3.11": [1344.0, 432.0],
+ "3.62": [1392.0, 384.0],
+ "3.75": [1440.0, 384.0],
+ "3.88": [1488.0, 384.0],
+ "4.0": [1536.0, 384.0],
+}
+
+ASPECT_RATIO_896 = {
+ "0.25": [448.0, 1792.0],
+ "0.26": [448.0, 1736.0],
+ "0.27": [448.0, 1680.0],
+ "0.28": [448.0, 1624.0],
+ "0.32": [504.0, 1568.0],
+ "0.33": [504.0, 1512.0],
+ "0.35": [504.0, 1456.0],
+ "0.4": [560.0, 1400.0],
+ "0.42": [560.0, 1344.0],
+ "0.48": [616.0, 1288.0],
+ "0.5": [616.0, 1232.0],
+ "0.52": [616.0, 1176.0],
+ "0.57": [672.0, 1176.0],
+ "0.6": [672.0, 1120.0],
+ "0.68": [728.0, 1064.0],
+ "0.72": [728.0, 1008.0],
+ "0.78": [784.0, 1008.0],
+ "0.82": [784.0, 952.0],
+ "0.88": [840.0, 952.0],
+ "0.94": [840.0, 896.0],
+ "1.0": [896.0, 896.0],
+ "1.07": [896.0, 840.0],
+ "1.13": [952.0, 840.0],
+ "1.21": [952.0, 784.0],
+ "1.29": [1008.0, 784.0],
+ "1.38": [1008.0, 728.0],
+ "1.46": [1064.0, 728.0],
+ "1.67": [1120.0, 672.0],
+ "1.75": [1176.0, 672.0],
+ "2.0": [1232.0, 616.0],
+ "2.09": [1288.0, 616.0],
+ "2.4": [1344.0, 560.0],
+ "2.5": [1400.0, 560.0],
+ "2.89": [1456.0, 504.0],
+ "3.0": [1512.0, 504.0],
+ "3.11": [1568.0, 504.0],
+ "3.62": [1624.0, 448.0],
+ "3.75": [1680.0, 448.0],
+ "3.88": [1736.0, 448.0],
+ "4.0": [1792.0, 448.0],
+}
+
+ASPECT_RATIO_1024 = {
+ "0.25": [512.0, 2048.0],
+ "0.26": [512.0, 1984.0],
+ "0.27": [512.0, 1920.0],
+ "0.28": [512.0, 1856.0],
+ "0.32": [576.0, 1792.0],
+ "0.33": [576.0, 1728.0],
+ "0.35": [576.0, 1664.0],
+ "0.4": [640.0, 1600.0],
+ "0.42": [640.0, 1536.0],
+ "0.48": [704.0, 1472.0],
+ "0.5": [704.0, 1408.0],
+ "0.52": [704.0, 1344.0],
+ "0.57": [768.0, 1344.0],
+ "0.6": [768.0, 1280.0],
+ "0.68": [832.0, 1216.0],
+ "0.72": [832.0, 1152.0],
+ "0.78": [896.0, 1152.0],
+ "0.82": [896.0, 1088.0],
+ "0.88": [960.0, 1088.0],
+ "0.94": [960.0, 1024.0],
+ "1.0": [1024.0, 1024.0],
+ "1.07": [1024.0, 960.0],
+ "1.13": [1088.0, 960.0],
+ "1.21": [1088.0, 896.0],
+ "1.29": [1152.0, 896.0],
+ "1.38": [1152.0, 832.0],
+ "1.46": [1216.0, 832.0],
+ "1.67": [1280.0, 768.0],
+ "1.75": [1344.0, 768.0],
+ "2.0": [1408.0, 704.0],
+ "2.09": [1472.0, 704.0],
+ "2.4": [1536.0, 640.0],
+ "2.5": [1600.0, 640.0],
+ "2.89": [1664.0, 576.0],
+ "3.0": [1728.0, 576.0],
+ "3.11": [1792.0, 576.0],
+ "3.62": [1856.0, 512.0],
+ "3.75": [1920.0, 512.0],
+ "3.88": [1984.0, 512.0],
+ "4.0": [2048.0, 512.0],
+}
+
+ASPECT_RATIO_1152 = {
+ "0.25": [576.0, 2304.0],
+ "0.26": [576.0, 2232.0],
+ "0.27": [576.0, 2160.0],
+ "0.28": [576.0, 2088.0],
+ "0.32": [648.0, 2016.0],
+ "0.33": [648.0, 1944.0],
+ "0.35": [648.0, 1872.0],
+ "0.4": [720.0, 1800.0],
+ "0.42": [720.0, 1728.0],
+ "0.48": [792.0, 1656.0],
+ "0.5": [792.0, 1584.0],
+ "0.52": [792.0, 1512.0],
+ "0.57": [864.0, 1512.0],
+ "0.6": [864.0, 1440.0],
+ "0.68": [936.0, 1368.0],
+ "0.72": [936.0, 1296.0],
+ "0.78": [1008.0, 1296.0],
+ "0.82": [1008.0, 1224.0],
+ "0.88": [1080.0, 1224.0],
+ "0.94": [1080.0, 1152.0],
+ "1.0": [1152.0, 1152.0],
+ "1.07": [1152.0, 1080.0],
+ "1.13": [1224.0, 1080.0],
+ "1.21": [1224.0, 1008.0],
+ "1.29": [1296.0, 1008.0],
+ "1.38": [1296.0, 936.0],
+ "1.46": [1368.0, 936.0],
+ "1.67": [1440.0, 864.0],
+ "1.75": [1512.0, 864.0],
+ "2.0": [1584.0, 792.0],
+ "2.09": [1656.0, 792.0],
+ "2.4": [1728.0, 720.0],
+ "2.5": [1800.0, 720.0],
+ "2.89": [1872.0, 648.0],
+ "3.0": [1944.0, 648.0],
+ "3.11": [2016.0, 648.0],
+ "3.62": [2088.0, 576.0],
+ "3.75": [2160.0, 576.0],
+ "3.88": [2232.0, 576.0],
+ "4.0": [2304.0, 576.0],
+}
+
+ASPECT_RATIO_1280 = {
+ "0.25": [640.0, 2560.0],
+ "0.26": [640.0, 2480.0],
+ "0.27": [640.0, 2400.0],
+ "0.28": [640.0, 2320.0],
+ "0.32": [720.0, 2240.0],
+ "0.33": [720.0, 2160.0],
+ "0.35": [720.0, 2080.0],
+ "0.4": [800.0, 2000.0],
+ "0.42": [800.0, 1920.0],
+ "0.48": [880.0, 1840.0],
+ "0.5": [880.0, 1760.0],
+ "0.52": [880.0, 1680.0],
+ "0.57": [960.0, 1680.0],
+ "0.6": [960.0, 1600.0],
+ "0.68": [1040.0, 1520.0],
+ "0.72": [1040.0, 1440.0],
+ "0.78": [1120.0, 1440.0],
+ "0.82": [1120.0, 1360.0],
+ "0.88": [1200.0, 1360.0],
+ "0.94": [1200.0, 1280.0],
+ "1.0": [1280.0, 1280.0],
+ "1.07": [1280.0, 1200.0],
+ "1.13": [1360.0, 1200.0],
+ "1.21": [1360.0, 1120.0],
+ "1.29": [1440.0, 1120.0],
+ "1.38": [1440.0, 1040.0],
+ "1.46": [1520.0, 1040.0],
+ "1.67": [1600.0, 960.0],
+ "1.75": [1680.0, 960.0],
+ "2.0": [1760.0, 880.0],
+ "2.09": [1840.0, 880.0],
+ "2.4": [1920.0, 800.0],
+ "2.5": [2000.0, 800.0],
+ "2.89": [2080.0, 720.0],
+ "3.0": [2160.0, 720.0],
+ "3.11": [2240.0, 720.0],
+ "3.62": [2320.0, 640.0],
+ "3.75": [2400.0, 640.0],
+ "3.88": [2480.0, 640.0],
+ "4.0": [2560.0, 640.0],
+}
+
+ASPECT_RATIO_1408 = {
+ "0.25": [704.0, 2816.0],
+ "0.26": [704.0, 2728.0],
+ "0.27": [704.0, 2640.0],
+ "0.28": [704.0, 2552.0],
+ "0.32": [792.0, 2464.0],
+ "0.33": [792.0, 2376.0],
+ "0.35": [792.0, 2288.0],
+ "0.4": [880.0, 2200.0],
+ "0.42": [880.0, 2112.0],
+ "0.48": [968.0, 2024.0],
+ "0.5": [968.0, 1936.0],
+ "0.52": [968.0, 1848.0],
+ "0.57": [1056.0, 1848.0],
+ "0.6": [1056.0, 1760.0],
+ "0.68": [1144.0, 1672.0],
+ "0.72": [1144.0, 1584.0],
+ "0.78": [1232.0, 1584.0],
+ "0.82": [1232.0, 1496.0],
+ "0.88": [1320.0, 1496.0],
+ "0.94": [1320.0, 1408.0],
+ "1.0": [1408.0, 1408.0],
+ "1.07": [1408.0, 1320.0],
+ "1.13": [1496.0, 1320.0],
+ "1.21": [1496.0, 1232.0],
+ "1.29": [1584.0, 1232.0],
+ "1.38": [1584.0, 1144.0],
+ "1.46": [1672.0, 1144.0],
+ "1.67": [1760.0, 1056.0],
+ "1.75": [1848.0, 1056.0],
+ "2.0": [1936.0, 968.0],
+ "2.09": [2024.0, 968.0],
+ "2.4": [2112.0, 880.0],
+ "2.5": [2200.0, 880.0],
+ "2.89": [2288.0, 792.0],
+ "3.0": [2376.0, 792.0],
+ "3.11": [2464.0, 792.0],
+ "3.62": [2552.0, 704.0],
+ "3.75": [2640.0, 704.0],
+ "3.88": [2728.0, 704.0],
+ "4.0": [2816.0, 704.0],
+}
+
+ASPECT_RATIO_2048 = {
+ "0.25": [1024.0, 4096.0],
+ "0.26": [1024.0, 3968.0],
+ "0.27": [1024.0, 3840.0],
+ "0.28": [1024.0, 3712.0],
+ "0.32": [1152.0, 3584.0],
+ "0.33": [1152.0, 3456.0],
+ "0.35": [1152.0, 3328.0],
+ "0.4": [1280.0, 3200.0],
+ "0.42": [1280.0, 3072.0],
+ "0.48": [1408.0, 2944.0],
+ "0.5": [1408.0, 2816.0],
+ "0.52": [1408.0, 2688.0],
+ "0.57": [1536.0, 2688.0],
+ "0.6": [1536.0, 2560.0],
+ "0.68": [1664.0, 2432.0],
+ "0.72": [1664.0, 2304.0],
+ "0.78": [1792.0, 2304.0],
+ "0.82": [1792.0, 2176.0],
+ "0.88": [1920.0, 2176.0],
+ "0.94": [1920.0, 2048.0],
+ "1.0": [2048.0, 2048.0],
+ "1.07": [2048.0, 1920.0],
+ "1.13": [2176.0, 1920.0],
+ "1.21": [2176.0, 1792.0],
+ "1.29": [2304.0, 1792.0],
+ "1.38": [2304.0, 1664.0],
+ "1.46": [2432.0, 1664.0],
+ "1.67": [2560.0, 1536.0],
+ "1.75": [2688.0, 1536.0],
+ "2.0": [2816.0, 1408.0],
+ "2.09": [2944.0, 1408.0],
+ "2.4": [3072.0, 1280.0],
+ "2.5": [3200.0, 1280.0],
+ "2.89": [3328.0, 1152.0],
+ "3.0": [3456.0, 1152.0],
+ "3.11": [3584.0, 1152.0],
+ "3.62": [3712.0, 1024.0],
+ "3.75": [3840.0, 1024.0],
+ "3.88": [3968.0, 1024.0],
+ "4.0": [4096.0, 1024.0],
+}
+
+ASPECT_RATIO_1440 = {
+ "0.25": [720.0, 2880.0],
+ "0.26": [720.0, 2790.0],
+ "0.27": [720.0, 2700.0],
+ "0.28": [720.0, 2610.0],
+ "0.32": [810.0, 2520.0],
+ "0.33": [810.0, 2430.0],
+ "0.35": [810.0, 2340.0],
+ "0.4": [900.0, 2250.0],
+ "0.42": [900.0, 2160.0],
+ "0.48": [990.0, 2070.0],
+ "0.5": [990.0, 1980.0],
+ "0.52": [990.0, 1890.0],
+ "0.57": [1080.0, 1890.0],
+ "0.6": [1080.0, 1800.0],
+ "0.68": [1170.0, 1710.0],
+ "0.72": [1170.0, 1620.0],
+ "0.78": [1260.0, 1620.0],
+ "0.82": [1260.0, 1530.0],
+ "0.88": [1350.0, 1530.0],
+ "0.94": [1350.0, 1440.0],
+ "1.0": [1440.0, 1440.0],
+ "1.07": [1440.0, 1350.0],
+ "1.13": [1530.0, 1350.0],
+ "1.21": [1530.0, 1260.0],
+ "1.29": [1620.0, 1260.0],
+ "1.38": [1620.0, 1170.0],
+ "1.46": [1710.0, 1170.0],
+ "1.67": [1800.0, 1080.0],
+ "1.75": [1890.0, 1080.0],
+ "2.0": [1980.0, 990.0],
+ "2.09": [2070.0, 990.0],
+ "2.4": [2160.0, 900.0],
+ "2.5": [2250.0, 900.0],
+ "2.89": [2340.0, 810.0],
+ "3.0": [2430.0, 810.0],
+ "3.11": [2520.0, 810.0],
+ "3.62": [2610.0, 720.0],
+ "3.75": [2700.0, 720.0],
+ "3.88": [2790.0, 720.0],
+ "4.0": [2880.0, 720.0],
+}
+
+ASPECT_RATIO_1536 = {
+ "0.25": [768.0, 3072.0],
+ "0.26": [768.0, 2976.0],
+ "0.27": [768.0, 2880.0],
+ "0.28": [768.0, 2784.0],
+ "0.32": [864.0, 2688.0],
+ "0.33": [864.0, 2592.0],
+ "0.35": [864.0, 2496.0],
+ "0.4": [960.0, 2400.0],
+ "0.42": [960.0, 2304.0],
+ "0.48": [1056.0, 2208.0],
+ "0.5": [1056.0, 2112.0],
+ "0.52": [1056.0, 2016.0],
+ "0.57": [1152.0, 2016.0],
+ "0.6": [1152.0, 1920.0],
+ "0.68": [1248.0, 1824.0],
+ "0.72": [1248.0, 1728.0],
+ "0.78": [1344.0, 1728.0],
+ "0.82": [1344.0, 1632.0],
+ "0.88": [1440.0, 1632.0],
+ "0.94": [1440.0, 1536.0],
+ "1.0": [1536.0, 1536.0],
+ "1.07": [1536.0, 1440.0],
+ "1.13": [1632.0, 1440.0],
+ "1.21": [1632.0, 1344.0],
+ "1.29": [1728.0, 1344.0],
+ "1.38": [1728.0, 1248.0],
+ "1.46": [1824.0, 1248.0],
+ "1.67": [1920.0, 1152.0],
+ "1.75": [2016.0, 1152.0],
+ "2.0": [2112.0, 1056.0],
+ "2.09": [2208.0, 1056.0],
+ "2.4": [2304.0, 960.0],
+ "2.5": [2400.0, 960.0],
+ "2.89": [2496.0, 864.0],
+ "3.0": [2592.0, 864.0],
+ "3.11": [2688.0, 864.0],
+ "3.62": [2784.0, 768.0],
+ "3.75": [2880.0, 768.0],
+ "3.88": [2976.0, 768.0],
+ "4.0": [3072.0, 768.0],
+}
+
+ASPECT_RATIO_1664 = {
+ "0.25": [832.0, 3328.0],
+ "0.26": [832.0, 3224.0],
+ "0.27": [832.0, 3120.0],
+ "0.28": [832.0, 3016.0],
+ "0.32": [936.0, 2912.0],
+ "0.33": [936.0, 2808.0],
+ "0.35": [936.0, 2704.0],
+ "0.4": [1040.0, 2600.0],
+ "0.42": [1040.0, 2496.0],
+ "0.48": [1144.0, 2392.0],
+ "0.5": [1144.0, 2288.0],
+ "0.52": [1144.0, 2184.0],
+ "0.57": [1248.0, 2184.0],
+ "0.6": [1248.0, 2080.0],
+ "0.68": [1352.0, 1976.0],
+ "0.72": [1352.0, 1872.0],
+ "0.78": [1456.0, 1872.0],
+ "0.82": [1456.0, 1768.0],
+ "0.88": [1560.0, 1768.0],
+ "0.94": [1560.0, 1664.0],
+ "1.0": [1664.0, 1664.0],
+ "1.07": [1664.0, 1560.0],
+ "1.13": [1768.0, 1560.0],
+ "1.21": [1768.0, 1456.0],
+ "1.29": [1872.0, 1456.0],
+ "1.38": [1872.0, 1352.0],
+ "1.46": [1976.0, 1352.0],
+ "1.67": [2080.0, 1248.0],
+ "1.75": [2184.0, 1248.0],
+ "2.0": [2288.0, 1144.0],
+ "2.09": [2392.0, 1144.0],
+ "2.4": [2496.0, 1040.0],
+ "2.5": [2600.0, 1040.0],
+ "2.89": [2704.0, 936.0],
+ "3.0": [2808.0, 936.0],
+ "3.11": [2912.0, 936.0],
+ "3.62": [3016.0, 832.0],
+ "3.75": [3120.0, 832.0],
+ "3.88": [3224.0, 832.0],
+ "4.0": [3328.0, 832.0],
+}
+
+ASPECT_RATIO_1792 = {
+ "0.25": [896.0, 3584.0],
+ "0.26": [896.0, 3472.0],
+ "0.27": [896.0, 3360.0],
+ "0.28": [896.0, 3248.0],
+ "0.32": [1008.0, 3136.0],
+ "0.33": [1008.0, 3024.0],
+ "0.35": [1008.0, 2912.0],
+ "0.4": [1120.0, 2800.0],
+ "0.42": [1120.0, 2688.0],
+ "0.48": [1232.0, 2576.0],
+ "0.5": [1232.0, 2464.0],
+ "0.52": [1232.0, 2352.0],
+ "0.57": [1344.0, 2352.0],
+ "0.6": [1344.0, 2240.0],
+ "0.68": [1456.0, 2128.0],
+ "0.72": [1456.0, 2016.0],
+ "0.78": [1568.0, 2016.0],
+ "0.82": [1568.0, 1904.0],
+ "0.88": [1680.0, 1904.0],
+ "0.94": [1680.0, 1792.0],
+ "1.0": [1792.0, 1792.0],
+ "1.07": [1792.0, 1680.0],
+ "1.13": [1904.0, 1680.0],
+ "1.21": [1904.0, 1568.0],
+ "1.29": [2016.0, 1568.0],
+ "1.38": [2016.0, 1456.0],
+ "1.46": [2128.0, 1456.0],
+ "1.67": [2240.0, 1344.0],
+ "1.75": [2352.0, 1344.0],
+ "2.0": [2464.0, 1232.0],
+ "2.09": [2576.0, 1232.0],
+ "2.4": [2688.0, 1120.0],
+ "2.5": [2800.0, 1120.0],
+ "2.89": [2912.0, 1008.0],
+ "3.0": [3024.0, 1008.0],
+ "3.11": [3136.0, 1008.0],
+ "3.62": [3248.0, 896.0],
+ "3.75": [3360.0, 896.0],
+ "3.88": [3472.0, 896.0],
+ "4.0": [3584.0, 896.0],
+}
+
+ASPECT_RATIO_1920 = {
+ "0.25": [960.0, 3840.0],
+ "0.26": [960.0, 3720.0],
+ "0.27": [960.0, 3600.0],
+ "0.28": [960.0, 3480.0],
+ "0.32": [1080.0, 3360.0],
+ "0.33": [1080.0, 3240.0],
+ "0.35": [1080.0, 3120.0],
+ "0.4": [1200.0, 3000.0],
+ "0.42": [1200.0, 2880.0],
+ "0.48": [1320.0, 2760.0],
+ "0.5": [1320.0, 2640.0],
+ "0.52": [1320.0, 2520.0],
+ "0.57": [1440.0, 2520.0],
+ "0.6": [1440.0, 2400.0],
+ "0.68": [1560.0, 2280.0],
+ "0.72": [1560.0, 2160.0],
+ "0.78": [1680.0, 2160.0],
+ "0.82": [1680.0, 2040.0],
+ "0.88": [1800.0, 2040.0],
+ "0.94": [1800.0, 1920.0],
+ "1.0": [1920.0, 1920.0],
+ "1.07": [1920.0, 1800.0],
+ "1.13": [2040.0, 1800.0],
+ "1.21": [2040.0, 1680.0],
+ "1.29": [2160.0, 1680.0],
+ "1.38": [2160.0, 1560.0],
+ "1.46": [2280.0, 1560.0],
+ "1.67": [2400.0, 1440.0],
+ "1.75": [2520.0, 1440.0],
+ "2.0": [2640.0, 1320.0],
+ "2.09": [2760.0, 1320.0],
+ "2.4": [2880.0, 1200.0],
+ "2.5": [3000.0, 1200.0],
+ "2.89": [3120.0, 1080.0],
+ "3.0": [3240.0, 1080.0],
+ "3.11": [3360.0, 1080.0],
+ "3.62": [3480.0, 960.0],
+ "3.75": [3600.0, 960.0],
+ "3.88": [3720.0, 960.0],
+ "4.0": [3840.0, 960.0],
+}
+
+ALL_ASPECT_RATIOS = [
+ ASPECT_RATIO_384,
+ ASPECT_RATIO_512,
+ ASPECT_RATIO_768,
+ ASPECT_RATIO_896,
+ ASPECT_RATIO_1024,
+ ASPECT_RATIO_1152,
+ ASPECT_RATIO_1280,
+ ASPECT_RATIO_1408,
+ # ASPECT_RATIO_1440,
+ ASPECT_RATIO_1536,
+ ASPECT_RATIO_1664,
+ ASPECT_RATIO_1792,
+ ASPECT_RATIO_1920,
+ ASPECT_RATIO_2048,
+]
+
+ALL_RESOLUTIONS = [384, 512, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048]
diff --git a/pipelines/vibe/edit_head.py b/pipelines/vibe/edit_head.py
new file mode 100644
index 000000000..5b96d3026
--- /dev/null
+++ b/pipelines/vibe/edit_head.py
@@ -0,0 +1,200 @@
+"""Meta Connector."""
+
+import torch
+from torch import nn
+from torch.nn import functional as func
+
+
+class Expert(nn.Module):
+ """Position-wise Feed-Forward Networks.
+
+ This consists of two linear transformations with a ReLU activation in between.
+ """
+
+ def __init__(self, d_model: int, d_expert: int) -> None:
+ """Initialize the Expert layer.
+
+ Args:
+ d_model (int): The dimension of the model.
+ d_expert (int): The dimension of the expert.
+ """
+ super().__init__()
+ self.d_model = d_model
+ self.d_expert = d_expert
+
+ # Linear transformation y = xW+b
+ self.fc1 = nn.Linear(self.d_model, self.d_expert, bias=True)
+ self.fc2 = nn.Linear(self.d_expert, self.d_model, bias=True)
+ self.fc3 = nn.Linear(self.d_model, self.d_expert, bias=True)
+
+ # Initialize weights
+ self._init_weights()
+
+ def _init_weights(self) -> None:
+ """Initialize the weights of the Expert layer."""
+ # init weights
+ nn.init.xavier_uniform_(self.fc1.weight)
+ nn.init.zeros_(self.fc1.bias)
+ nn.init.xavier_uniform_(self.fc2.weight)
+ nn.init.zeros_(self.fc2.bias)
+ nn.init.xavier_uniform_(self.fc3.weight)
+ nn.init.zeros_(self.fc3.bias)
+
+ def forward(self, embedding: torch.Tensor) -> torch.Tensor:
+ """Forward pass of the Expert layer.
+
+ Args:
+ embedding (torch.Tensor): The input tensor.
+
+ Returns:
+ torch.Tensor: The output tensor.
+ """
+ inner_emb = func.silu(self.fc1(embedding)) * self.fc3(embedding)
+ return self.fc2(inner_emb)
+
+
+class MetaAttentionBlock(nn.Module):
+ """Meta Attention Block."""
+
+ def __init__(self, d_model: int, n_heads: int, dropout: float = 0.05) -> None:
+ """Initialize the Meta Attention Block.
+
+ Args:
+ d_model (int): The dimension of the embeddings.
+ n_heads (int): The number of attention heads.
+ dropout (float): The dropout rate.
+ """
+ super().__init__()
+ self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
+
+ # Dropout layers for residual connections
+ self.dropout_attn = nn.Dropout(dropout)
+ self.dropout_ffn = nn.Dropout(dropout)
+
+ self.mlp = Expert(d_model, d_model * 4)
+
+ self.input_norm = nn.LayerNorm(d_model)
+ self.output_norm = nn.LayerNorm(d_model)
+
+ # Initialize weights
+ self._init_weights()
+
+ def _init_weights(self) -> None:
+ """Initialize weights for attention and MLP blocks."""
+ # Initialize MultiheadAttention weights
+ # The in_proj_weight contains Q, K, V concatenated
+ nn.init.xavier_uniform_(self.attn.in_proj_weight)
+ if self.attn.in_proj_bias is not None:
+ nn.init.zeros_(self.attn.in_proj_bias)
+
+ # Initialize the output projection of MultiheadAttention
+ nn.init.xavier_uniform_(self.attn.out_proj.weight)
+ if self.attn.out_proj.bias is not None:
+ nn.init.zeros_(self.attn.out_proj.bias)
+
+ nn.init.ones_(self.input_norm.weight)
+ nn.init.zeros_(self.input_norm.bias)
+ nn.init.ones_(self.output_norm.weight)
+ nn.init.zeros_(self.output_norm.bias)
+
+ def forward(self, embedding: torch.Tensor) -> torch.Tensor:
+ """Forward pass of the MetaAttentionBlock.
+
+ Args:
+ embedding (torch.Tensor): The input embeddings. Shape: (batch_size, seq_len, hidden_dim).
+
+ Returns:
+ torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim).
+ """
+ # apply layer normalization to embeddings
+ normed_embedding = self.input_norm(embedding)
+
+ # apply attention and skip connection
+ attn_output, _ = self.attn(normed_embedding, normed_embedding, normed_embedding, need_weights=False)
+ embedding = embedding + self.dropout_attn(attn_output)
+
+ output = self.mlp(self.output_norm(embedding))
+ return embedding + self.dropout_ffn(output)
+
+
+class MetaConnector(nn.Module):
+ """Meta Connector."""
+
+ def __init__(
+ self,
+ input_dim: int = 4096,
+ hidden_dim: int = 2048,
+ output_dim: int = 768,
+ n_stacks: int = 8,
+ heads: int = 16,
+ *,
+ use_input_proj: bool = True,
+ ) -> None:
+ """Initialize the MetaConnector.
+
+ Args:
+ input_dim (int): The dimension of the input.
+ hidden_dim (int): The dimension of the hidden state.
+ output_dim (int): The dimension of the output.
+ n_stacks (int): The number of attention layers.
+ heads (int): The number of attention heads.
+ use_input_proj (bool): Whether to use input projection.
+ """
+ super().__init__()
+ self.output_dim = output_dim
+ self.input_dim = input_dim
+ self.use_input_proj = use_input_proj
+
+ # Input Projection
+ if use_input_proj:
+ self.input_proj = nn.Linear(input_dim, hidden_dim)
+ else:
+ hidden_dim = input_dim
+
+ # Attention blocks
+ self.attn_blocks = nn.Sequential(*[MetaAttentionBlock(hidden_dim, heads) for _ in range(n_stacks)])
+
+ # Output Projection
+ self.output_proj = nn.Sequential(nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, output_dim))
+
+ # Initialize weights
+ self._init_weights()
+
+ def _init_weights(self) -> None:
+ """Initialize weights."""
+ # Initialize input_proj
+ if self.use_input_proj:
+ nn.init.xavier_uniform_(self.input_proj.weight)
+ if self.input_proj.bias is not None:
+ nn.init.zeros_(self.input_proj.bias)
+
+ # Initialize output_proj
+ for module in self.output_proj:
+ if isinstance(module, nn.Linear):
+ nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.zeros_(module.bias)
+ elif isinstance(module, nn.LayerNorm):
+ nn.init.ones_(module.weight)
+ nn.init.zeros_(module.bias)
+
+ def forward(self, hidden_states: torch.Tensor, meta_queries: torch.Tensor) -> dict[str, torch.Tensor]:
+ """Forward pass of the Meta Connector.
+
+ Args:
+ hidden_states (torch.Tensor): The input hidden states. Shape: (batch_size, seq_len, hidden_dim).
+ meta_queries (torch.Tensor): The input meta queries. Shape: (batch_size, seq_len, hidden_dim).
+
+ Returns:
+ torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim).
+ """
+ hidden_states = hidden_states + meta_queries
+ if self.use_input_proj:
+ hidden_states = self.input_proj(hidden_states)
+
+ # apply attention blocks
+ for attn_block in self.attn_blocks:
+ hidden_states = attn_block(hidden_states)
+
+ # apply output projection
+ return self.output_proj(hidden_states)
diff --git a/pipelines/vibe/vibe_sana_editing.py b/pipelines/vibe/vibe_sana_editing.py
new file mode 100644
index 000000000..3b88bed2e
--- /dev/null
+++ b/pipelines/vibe/vibe_sana_editing.py
@@ -0,0 +1,624 @@
+"""Transformer Sana model."""
+
+from typing import Any
+
+import torch
+from diffusers import ModelMixin, SanaTransformer2DModel
+from diffusers.configuration_utils import register_to_config
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import PatchEmbed, PixArtAlphaTextProjection
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormSingle, RMSNorm
+from diffusers.models.transformers.sana_transformer import (
+ GLUMBConv,
+ SanaAttnProcessor2_0,
+ SanaCombinedTimestepGuidanceEmbeddings,
+ SanaModulatedNorm,
+)
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from torch import nn
+
+from .edit_head import MetaConnector
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class SanaLinearAttnProcessor2_0:
+ r"""Processor for implementing scaled dot-product linear attention."""
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Execute the forward pass of the SanaLinearAttnProcessor2_0.
+
+ We need that redifinition because default implementation do not take into account the attention mask.
+
+ Args:
+ attn (Attention): The attention object.
+ hidden_states (torch.Tensor): The hidden states.
+ encoder_hidden_states (torch.Tensor | None): The encoder hidden states. In case of cross attention.
+ attention_mask (torch.Tensor | None): The attention mask.
+
+ Returns:
+ torch.Tensor: The processed hidden states.
+ """
+ original_dtype = hidden_states.dtype
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ query = query.transpose(1, 2).unflatten(1, (attn.heads, -1))
+ key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
+ value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
+
+ query = nn.functional.relu(query)
+ key = nn.functional.relu(key)
+
+ query, key, value = query.float(), key.float(), value.float()
+
+ if attention_mask is not None:
+ # attention_mask([B, S], dtype=torch.bool): True - keep, False - ignore.
+ # transform to (1.0 - keep, 0.0 - ignore).
+ mask = attention_mask if attention_mask.dtype == torch.bool else attention_mask > -1
+ mask = mask.to(key.dtype)
+ if mask.ndim == 2:
+ mask = mask.unsqueeze(1) # -> [B, 1, S]
+
+ # Null masked tokens in key and value before torch.matmul operation
+ key = key * mask.unsqueeze(-1) # (B, _, S, _) * (B, 1, S, 1)
+ value = value * mask.unsqueeze(2) # (B, _, _, S) * (B, 1, 1, S)
+
+ value = nn.functional.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
+ scores = torch.matmul(value, key)
+ hidden_states = torch.matmul(scores, query)
+
+ hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + 1e-15)
+ hidden_states = hidden_states.flatten(1, 2).transpose(1, 2)
+ hidden_states = hidden_states.to(original_dtype)
+
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if original_dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return hidden_states
+
+
+class SanaTransformerBlock(nn.Module):
+ """Transformer block introduced in [Sana](huggingface.co/papers/2410.10629)."""
+
+ def __init__(
+ self,
+ dim: int = 2240,
+ num_attention_heads: int = 70,
+ attention_head_dim: int = 32,
+ dropout: float = 0.0,
+ num_cross_attention_heads: int = 20,
+ cross_attention_head_dim: int = 112,
+ cross_attention_dim: int = 2240,
+ norm_eps: float = 1e-6,
+ mlp_ratio: float = 2.5,
+ qk_norm: str | None = None,
+ *,
+ attention_out_bias: bool = True,
+ attention_bias: bool = True,
+ norm_elementwise_affine: bool = False,
+ ) -> None:
+ """Initialize the SanaTransformerBlock.
+
+ Args:
+ dim (int): The dimension of the input/output hidden states.
+ num_attention_heads (int): The number of heads for self-attention.
+ attention_head_dim (int): The dimension of each self-attention head.
+ dropout (float): The dropout probability for attention layers.
+ num_cross_attention_heads (int): The number of heads for cross-attention.
+ cross_attention_head_dim (int): The dimension of each cross-attention head.
+ cross_attention_dim (int): The dimension of the cross-attention context.
+ norm_eps (float): Epsilon value for layer normalization stability.
+ mlp_ratio (float): Expansion ratio for the feed-forward network hidden dimension.
+ qk_norm (str | None): Normalization method for Query/Key vectors.
+ attention_out_bias (bool): Whether to include bias in the attention output projection.
+ attention_bias (bool): Whether to include bias in the attention Q/K/V projections.
+ norm_elementwise_affine (bool): Whether to learn affine parameters for normalization.
+ """
+ super().__init__()
+
+ # 1. Self Attention
+ self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps)
+ self.attn1 = Attention(
+ query_dim=dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ kv_heads=num_attention_heads if qk_norm is not None else None,
+ qk_norm=qk_norm,
+ dropout=dropout,
+ bias=attention_bias,
+ cross_attention_dim=None,
+ processor=SanaLinearAttnProcessor2_0(), # type: ignore[arg-type]
+ )
+
+ # 2. Cross Attention
+ if cross_attention_dim is not None:
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+ self.attn2 = Attention(
+ query_dim=dim,
+ qk_norm=qk_norm,
+ kv_heads=num_cross_attention_heads if qk_norm is not None else None,
+ cross_attention_dim=cross_attention_dim,
+ heads=num_cross_attention_heads,
+ dim_head=cross_attention_head_dim,
+ dropout=dropout,
+ bias=True,
+ out_bias=attention_out_bias,
+ processor=SanaAttnProcessor2_0(), # type: ignore[arg-type]
+ )
+
+ # 3. Feed-forward
+ self.ff = GLUMBConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False)
+
+ self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ encoder_hidden_states: torch.Tensor | None,
+ encoder_attention_mask: torch.Tensor | None,
+ timestep: torch.LongTensor,
+ height: int,
+ width: int,
+ ) -> torch.Tensor:
+ """Executes the forward pass of the transformer block.
+
+ Args:
+ hidden_states (torch.Tensor): Input hidden states of shape `(batch_size, seq_len, dim)`.
+ attention_mask (torch.Tensor | None): Mask for self-attention.
+ encoder_hidden_states (torch.Tensor | None): Conditional embeddings (e.g., text) for cross-attention.
+ encoder_attention_mask (torch.Tensor | None): Mask for cross-attention.
+ timestep (torch.LongTensor): Timestep embeddings used for adaptive modulation of normalization layers.
+ height (int): Spatial height of the latent image (used for reshaping in FFN).
+ width (int): Spatial width of the latent image (used for reshaping in FFN).
+
+ Returns:
+ torch.Tensor: Transformed hidden states of shape `(batch_size, seq_len, dim)`.
+ """
+ batch_size = hidden_states.shape[0]
+
+ # 1. Modulation
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+ self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+ ).chunk(6, dim=1)
+
+ # 2. Self Attention
+ norm_hidden_states = self.norm1(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+ norm_hidden_states = norm_hidden_states.to(hidden_states.dtype)
+
+ attn_output = self.attn1(norm_hidden_states, attention_mask=attention_mask)
+ hidden_states = hidden_states + gate_msa * attn_output
+
+ # 3. Cross Attention
+ if self.attn2 is not None:
+ attn_output = self.attn2(
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ )
+ hidden_states = attn_output + hidden_states
+
+ # 4. Feed-forward
+ norm_hidden_states = self.norm2(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+ if width * height == norm_hidden_states.shape[1]:
+ norm_hidden_states = norm_hidden_states.unflatten(1, (height, width)).permute(0, 3, 1, 2)
+ ff_output = self.ff(norm_hidden_states)
+ ff_output = ff_output.flatten(2, 3).permute(0, 2, 1)
+ else:
+ # "seq_cat" case
+ # apply the same local processing to each of the parts independently of each other
+ # global interaction has already occurred, work with spatial features separately
+ norm_hidden_states_1, norm_hidden_states_2 = torch.chunk(norm_hidden_states, 2, dim=1)
+
+ hidden_states_reshaped_1 = norm_hidden_states_1.unflatten(1, (height, width)).permute(0, 3, 1, 2)
+ hidden_states_reshaped_2 = norm_hidden_states_2.unflatten(1, (height, width)).permute(0, 3, 1, 2)
+
+ ff_output_1 = self.ff(hidden_states_reshaped_1).flatten(2, 3).permute(0, 2, 1)
+ ff_output_2 = self.ff(hidden_states_reshaped_2).flatten(2, 3).permute(0, 2, 1)
+
+ ff_output = torch.cat([ff_output_1, ff_output_2], dim=1)
+
+ return hidden_states + gate_mlp * ff_output
+
+
+class VIBESanaEditingModel(SanaTransformer2DModel):
+ """VIBE Sana editing model specialized for instruction-based image editing.
+
+ This model extends `SanaTransformer2DModel` to support advanced editing capabilities.
+ """
+
+ @register_to_config
+ def __init__(
+ self,
+ in_channels: int = 32,
+ out_channels: int | None = 32,
+ num_attention_heads: int = 70,
+ attention_head_dim: int = 32,
+ num_layers: int = 20,
+ num_cross_attention_heads: int = 20,
+ cross_attention_head_dim: int = 112,
+ cross_attention_dim: int = 2240,
+ caption_channels: int = 2304,
+ mlp_ratio: float = 2.5,
+ dropout: float = 0.0,
+ sample_size: int = 32,
+ patch_size: int = 1,
+ norm_eps: float = 1e-6,
+ interpolation_scale: int | None = None,
+ qk_norm: str | None = None,
+ timestep_scale: float = 1.0,
+ input_condition_type: str = "channel_cat",
+ edit_head_input_dim: int = 2048,
+ edit_head_stacks_num: int = 4,
+ num_meta_queries: int = 224,
+ meta_queries_dim: int = 2048,
+ *,
+ attention_bias: bool = False,
+ norm_elementwise_affine: bool = False,
+ guidance_embeds: bool = False,
+ ) -> None:
+ """Initialize the VIBESanaEditingModel.
+
+ Args:
+ in_channels (int): Number of channels in the input latent representation.
+ out_channels (int | None): Number of output channels. Defaults to `in_channels`.
+ num_attention_heads (int): Number of attention heads in transformer blocks.
+ attention_head_dim (int): Dimension of each attention head.
+ num_layers (int): Number of transformer blocks in the model.
+ num_cross_attention_heads (int): Number of heads for cross-attention layers.
+ cross_attention_head_dim (int): Dimension of each cross-attention head.
+ cross_attention_dim (int): Dimension of the cross-attention context embedding.
+ caption_channels (int): Dimension of the input caption/text embeddings.
+ mlp_ratio (float): Expansion ratio for the feed-forward networks.
+ dropout (float): Dropout probability.
+ sample_size (int): Spatial size of the input sample (height/width).
+ patch_size (int): Size of the patches extracted from the input latent.
+ norm_eps (float): Epsilon for layer normalization.
+ interpolation_scale (int | None): Scale factor for positional embedding interpolation.
+ qk_norm (str | None): Normalization type for Query/Key (e.g., 'rms_norm').
+ timestep_scale (float): Scale factor for the timestep.
+ input_condition_type (str): Method for conditioning on the input image. Options: 'channel_cat', 'seq_cat'.
+ edit_head_input_dim (int): Input dimension for the `MetaConnector` edit head.
+ edit_head_stacks_num (int): Number of layers in the `MetaConnector`.
+ num_meta_queries (int): Number of meta tokens to extract from the text embeddings for the edit head.
+ meta_queries_dim (int): Dimension of the meta query embeddings.
+ attention_bias (bool): Whether to use bias in attention projections.
+ norm_elementwise_affine (bool): Whether to learn affine parameters in normalization.
+ guidance_embeds (bool): Whether to use additional guidance embeddings.
+ """
+ super(SanaTransformer2DModel, self).__init__()
+
+ out_channels = out_channels or in_channels
+ inner_dim = num_attention_heads * attention_head_dim
+
+ # 1. Patch Embedding
+ self.patch_embed = PatchEmbed(
+ height=sample_size,
+ width=sample_size,
+ patch_size=patch_size,
+ in_channels=in_channels,
+ embed_dim=inner_dim,
+ interpolation_scale=interpolation_scale,
+ pos_embed_type="sincos" if interpolation_scale is not None else None,
+ )
+
+ # 2. Additional condition embeddings
+ if guidance_embeds:
+ self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim)
+ else:
+ self.time_embed = AdaLayerNormSingle(inner_dim) # type: ignore[assignment]
+
+ self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
+ self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True)
+
+ # 3. Transformer blocks
+ self.transformer_blocks = nn.ModuleList(
+ [
+ SanaTransformerBlock(
+ dim=inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ dropout=dropout,
+ num_cross_attention_heads=num_cross_attention_heads,
+ cross_attention_head_dim=cross_attention_head_dim,
+ cross_attention_dim=cross_attention_dim,
+ attention_bias=attention_bias,
+ norm_elementwise_affine=norm_elementwise_affine,
+ norm_eps=norm_eps,
+ mlp_ratio=mlp_ratio,
+ qk_norm=qk_norm,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ # 4. Output blocks
+ self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5)
+ self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
+
+ # 5. Meta queries embedding
+ self.num_meta_queries = num_meta_queries
+ self.meta_queries = nn.Parameter(torch.randn(num_meta_queries, meta_queries_dim) / meta_queries_dim**0.5)
+
+ # 6. Edit head
+ self.input_condition_type = input_condition_type
+ self.edit_head = MetaConnector(
+ input_dim=edit_head_input_dim,
+ n_stacks=edit_head_stacks_num,
+ output_dim=caption_channels,
+ )
+ self.gradient_checkpointing = False
+
+ def get_hidden_states_for_meta_tokens(
+ self,
+ text_encoder_hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ ) -> torch.Tensor:
+ """Extracts hidden states corresponding to meta tokens from the text encoder output.
+
+ This method identifies the positions of meta tokens in the input sequence using
+ the attention mask and extracts their embeddings. It assumes the meta tokens
+ are located at the end of the valid sequence.
+
+ Args:
+ text_encoder_hidden_states (torch.Tensor): Output embeddings from the text encoder
+ of shape `(batch_size, seq_len, hidden_dim)`.
+ attention_mask (torch.Tensor): Boolean mask indicating valid tokens (True) and
+ padding (False), used to locate the end of the prompt.
+
+ Returns:
+ torch.Tensor: Extracted meta token embeddings of shape `(batch_size, num_meta_queries, hidden_dim)`.
+ """
+ img_hidden_states: list[torch.Tensor] = []
+ for i in range(text_encoder_hidden_states.shape[0]):
+ hidden_states_masked = text_encoder_hidden_states[i][attention_mask[i]]
+ img_hidden_states.append(hidden_states_masked[-self.num_meta_queries :, :].unsqueeze(0))
+ return torch.cat(img_hidden_states, dim=0)
+
+ def forward_edit_heads(
+ self,
+ batch_size: int,
+ meta_tokens_hidden_states: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Processes meta token embeddings through the edit head for both conditional and unconditional generation.
+
+ This method generates editing embeddings for:
+ 1. **Conditional input** (`hid_edit`): Using the extracted `meta_tokens_hidden_states`.
+ 2. **Unconditional/Null input** (`hid_null`): Using zeroed hidden states, facilitating
+ Classifier-Free Guidance (CFG) during inference.
+
+ Args:
+ batch_size (int): The batch size of the generation.
+ meta_tokens_hidden_states (torch.Tensor): Input hidden states for the meta tokens
+ extracted from the text prompt.
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+ - `hid_edit` (torch.Tensor): Editing embeddings for the conditional branch.
+ - `hid_null` (torch.Tensor): Editing embeddings for the unconditional (null) branch.
+ """
+ # repeat meta tokens embeddings to fit batch_size
+ meta_tokens_embeddings = self.meta_queries.unsqueeze(dim=0)
+ meta_tokens_embeddings = meta_tokens_embeddings.repeat(batch_size, 1, 1)
+
+ # create zero hidden states for null prompts
+ zero_hidden_states = torch.zeros_like(meta_tokens_hidden_states)
+
+ # get hidden states for real prompts
+ hid_edit = self.edit_head(meta_tokens_hidden_states, meta_tokens_embeddings)
+
+ # get hidden states for null prompts
+ hid_null = self.edit_head(zero_hidden_states, meta_tokens_embeddings)
+
+ return hid_edit, hid_null
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ guidance: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ controlnet_block_samples: tuple[torch.Tensor] | None = None,
+ *,
+ t2i_samples: list[bool] | None = None,
+ return_dict: bool = True,
+ ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
+ """Execute the forward pass of the VIBE Sana Editing Model.
+
+ Args:
+ hidden_states (torch.Tensor): Noisy latent image input of shape
+ `(batch_size, num_channels, height, width)`.
+ encoder_hidden_states (torch.Tensor): Conditional text embeddings of shape
+ `(batch_size, seq_len, hidden_dim)`.
+ timestep (torch.Tensor): Denoising timesteps.
+ guidance (torch.Tensor | None, optional): Guidance embeddings for the diffusion process.
+ encoder_attention_mask (torch.Tensor | None, optional): Mask for `encoder_hidden_states`.
+ attention_mask (torch.Tensor | None, optional): Mask for `hidden_states`.
+ attention_kwargs (dict[str, Any] | None, optional): Extra arguments for attention processors.
+ controlnet_block_samples (tuple[torch.Tensor] | None, optional): Additional residuals
+ from a ControlNet model, if used.
+ t2i_samples (list[bool] | None, optional): Marker if the sample is a t2i sample.
+ return_dict (bool): If True, returns a `Transformer2DModelOutput`. Otherwise, returns a tuple.
+
+ Returns:
+ tuple[torch.Tensor, ...] | Transformer2DModelOutput: The denoised latent output.
+ If `return_dict` is True, returns `Transformer2DModelOutput(sample=output)`.
+
+ Raises:
+ ValueError: If `input_condition_type` is not 'channel_cat' or 'seq_cat'.
+ """
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ scale_lora_layers(self, lora_scale)
+ elif attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+ logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
+
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+ # ensure attention_mask is a bias
+ # assume that mask is expressed as:
+ # # (1 = keep, 0 = discard)
+ # # convert mask into a bias that can be added to attention scores:
+ # # (keep = +0, discard = -10000.0)
+ encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1) # type: ignore[union-attr]
+
+ # 1. Input
+ batch_size, num_channels, height, width = hidden_states.shape
+
+ if self.input_condition_type == "channel_cat": # type: ignore
+ # For mixed batch of t2i and editing samples, we need to mask the input image embeddings for the t2i samples
+ if t2i_samples is not None and any(t2i_samples):
+ hidden_states, input_img_embeds = hidden_states.split(num_channels // 2, dim=1) # type: ignore
+ mask = torch.tensor(t2i_samples, dtype=torch.bool, device=input_img_embeds.device).view(-1, 1, 1, 1)
+ input_img_embeds = torch.where(mask, torch.zeros_like(input_img_embeds), input_img_embeds)
+ hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1) # type: ignore
+ hidden_states = self.patch_embed(hidden_states)
+ attention_mask = None # There is no need for attention mask in channel_cat case
+
+ elif self.input_condition_type == "seq_cat":
+ num_channels = hidden_states.shape[1]
+ hidden_states, input_img_embeds = hidden_states.split(num_channels // 2, dim=1)
+ hidden_states = self.patch_embed(hidden_states)
+
+ if t2i_samples is not None and all(t2i_samples):
+ # All samples are t2i samples, so there is no need for attention mask
+ # Also, we don't need to patch the input image embeddings
+ attention_mask = None
+ elif t2i_samples is None or (t2i_samples is not None and not any(t2i_samples)):
+ # All samples are editing samples, still no need for attention mask
+ # But we need to patch the input image embeddings
+ input_img_embeds = self.patch_embed(input_img_embeds)
+ hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1)
+ attention_mask = None
+ else:
+ # Mixed samples.
+ input_img_embeds = self.patch_embed(input_img_embeds)
+ # Create attnention masks for hidden_states and input_img_embeds
+ hidden_stated_mask = torch.ones(
+ hidden_states.shape[0],
+ hidden_states.shape[1],
+ dtype=torch.bool,
+ device=hidden_states.device,
+ ) # Shape: [bs, seq_len]
+ input_img_embeds_mask = torch.ones(
+ input_img_embeds.shape[0],
+ input_img_embeds.shape[1],
+ dtype=torch.bool,
+ device=input_img_embeds.device,
+ ) # Shape: [bs, seq_len]
+
+ # Set False for t2i samples in input_img_embeds_mask
+ if t2i_samples is not None:
+ t2i_samples_tensor = torch.as_tensor(t2i_samples, dtype=torch.bool, device=hidden_states.device)
+ input_img_embeds_mask[t2i_samples_tensor] = False
+
+ hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1) # type: ignore
+ attention_mask = torch.cat([hidden_stated_mask, input_img_embeds_mask], dim=1)
+ else:
+ msg = f"Invalid input condition type: {self.input_condition_type}"
+ raise ValueError(msg) # type: ignore
+
+ patch_size = self.config.patch_size # type: ignore[attr-defined]
+ post_patch_height, post_patch_width = height // patch_size, width // patch_size
+
+ if guidance is not None:
+ timestep, embedded_timestep = self.time_embed(timestep, guidance=guidance, hidden_dtype=hidden_states.dtype)
+ else:
+ timestep, embedded_timestep = self.time_embed(
+ timestep,
+ batch_size=batch_size,
+ hidden_dtype=hidden_states.dtype,
+ )
+
+ encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
+ encoder_hidden_states = self.caption_norm(encoder_hidden_states)
+
+ # 2. Transformer blocks
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ for index_block, block in enumerate(self.transformer_blocks):
+ hidden_states = self._gradient_checkpointing_func( # type: ignore[attr-defined]
+ block,
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ timestep,
+ post_patch_height,
+ post_patch_width,
+ )
+ if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples):
+ hidden_states = hidden_states + controlnet_block_samples[index_block - 1]
+
+ else:
+ for index_block, block in enumerate(self.transformer_blocks):
+ hidden_states = block(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ timestep,
+ post_patch_height,
+ post_patch_width,
+ )
+ if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples):
+ hidden_states = hidden_states + controlnet_block_samples[index_block - 1]
+
+ # 3. Normalization
+ hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table)
+ hidden_states = self.proj_out(hidden_states)
+
+ # 5. Unpatchify
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ post_patch_height,
+ post_patch_width,
+ patch_size,
+ patch_size,
+ -1,
+ )
+ hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
+ output = hidden_states.reshape(batch_size, -1, post_patch_height * patch_size, post_patch_width * patch_size)
+
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+
+ return Transformer2DModelOutput(sample=output)
diff --git a/pipelines/vibe/vibe_sana_pipeline.py b/pipelines/vibe/vibe_sana_pipeline.py
new file mode 100644
index 000000000..d4ac38961
--- /dev/null
+++ b/pipelines/vibe/vibe_sana_pipeline.py
@@ -0,0 +1,856 @@
+"""VIBE Sana Editing pipeline."""
+
+from typing import Any
+
+import torch
+from diffusers.image_processor import PixArtImageProcessor
+from diffusers.models import AutoencoderDC
+from diffusers.pipelines.sana.pipeline_output import SanaPipelineOutput
+from diffusers.pipelines.sana.pipeline_sana import SanaPipeline, retrieve_timesteps
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import is_torch_xla_available
+from PIL import Image
+from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+from .aspects_multiscale import ASPECT_RATIO_512, ASPECT_RATIO_1024, ASPECT_RATIO_2048
+from .vibe_sana_editing import VIBESanaEditingModel
+
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm # type: ignore # pylint: disable
+
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+
+class VIBESanaEditingPipeline(SanaPipeline):
+ """Image Editing Pipeline for Sana."""
+
+ edit_query_template = "what will this image be like if {instruction}"
+ t2i_query_template = "generate the image by description: {instruction}"
+ min_pixels = 352 * 352 # min size of image in Qwen3VLProcessor
+ max_pixels = 672 * 672 # max size of image in Qwen3VLProcessor
+ bin_512_space = 512 * 512
+ bin_1024_space = 1024 * 1024
+ bin_2048_space = 2048 * 2048
+
+ def __init__(
+ self,
+ tokenizer: Qwen3VLProcessor,
+ text_encoder: Qwen3VLForConditionalGeneration,
+ vae: AutoencoderDC,
+ transformer: VIBESanaEditingModel,
+ scheduler: DPMSolverMultistepScheduler,
+ ) -> None:
+ """Initialize the SanaIP2P pipeline.
+
+ Args:
+ tokenizer (Qwen3VLProcessor): The tokenizer for the text encoder.
+ text_encoder (Qwen3VLForConditionalGeneration): The text encoder for the editing prompt.
+ vae (AutoencoderDC): The VAE model.
+ transformer (VIBESanaEditingModel): The editing transformer.
+ scheduler (DPMSolverMultistepScheduler): The sana scheduler.
+ """
+ super(SanaPipeline, self).__init__()
+
+ self.register_modules(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer,
+ scheduler=scheduler,
+ )
+
+ self.vae_scale_factor = (
+ 2 ** (len(self.vae.config.encoder_block_out_channels) - 1)
+ if hasattr(self, "vae") and self.vae is not None
+ else 32
+ )
+ self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+ @property
+ def do_classifier_free_guidance(self) -> bool:
+ """Property to check if classifier free guidance is enabled.
+
+ Returns:
+ bool: True if classifier free guidance is enabled, False otherwise.
+ """
+ return self._guidance_scale > 1
+
+ @property
+ def do_image_guidance(self) -> bool:
+ """Property to check if image guidance is enabled.
+
+ Returns:
+ bool: True if image guidance is enabled, False otherwise.
+ """
+ return self._image_guidance_scale >= 1
+
+ def get_closest_size(self, height: int, width: int) -> tuple[int, int]:
+ """Get the closest size to the original size.
+
+ Args:
+ height (int): The height of the image.
+ width (int): The width of the image.
+
+ Returns:
+ tuple[int, int]: The closest size to the original size.
+ """
+ sample_area = height * width
+ if sample_area <= self.bin_512_space * 2 + 256**2:
+ aspect_ratio_bin = ASPECT_RATIO_512
+ elif sample_area <= self.bin_1024_space * 2 + 512**2:
+ aspect_ratio_bin = ASPECT_RATIO_1024
+ else:
+ aspect_ratio_bin = ASPECT_RATIO_2048
+ height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
+ return height, width
+
+ def prepare_inputs_batch(
+ self,
+ instruction: str | list[str],
+ image: Image.Image | None = None,
+ ) -> dict[str, torch.Tensor]:
+ """Prepare inputs for the batch.
+
+ Args:
+ instruction (str | list[str]): Edit prompts.
+ image (Image.Image | None): The input image or None if t2i generation is enabled.
+
+ Returns:
+ dict[str, torch.Tensor]: The inputs for the batch.
+ - input_ids: The input IDs.
+ - attention_mask: The attention mask.
+ - image_grid_thw: The image grid size.
+ - pixel_values: The pixel values.
+ """
+ # Prepare instructions
+ instructions = [instruction] if isinstance(instruction, str) else instruction
+ samples = []
+ for instr in instructions:
+ query = (
+ self.edit_query_template.format(instruction=instr.strip())
+ if image is not None
+ else self.t2i_query_template.format(instruction=instr.strip())
+ )
+ user_content = [{"type": "image", "image": image}] if image is not None else []
+ user_content.append({"type": "text", "text": query})
+ message = [{"role": "user", "content": user_content}]
+ text = self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=False)
+ sample = {"text": text}
+ if image is not None:
+ sample["image"] = image
+ samples.append(sample)
+
+ # Invoke the processor to get the input ids and attention mask
+ texts = [sample["text"] for sample in samples]
+ images = [sample["image"] for sample in samples] if image is not None else None
+ kwargs = {"min_pixels": self.min_pixels, "max_pixels": self.max_pixels}
+ inputs = self.tokenizer(text=texts, images=images, padding=True, return_tensors="pt", **kwargs).data
+ return {key: value.to(self._execution_device) for key, value in inputs.items()}
+
+ def prepare_inputs_for_meta_queries(
+ self,
+ inputs_embeds: torch.Tensor,
+ attention_mask: torch.Tensor,
+ image_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+ """Add meta queries into input embeddings before any padding tokens.
+
+ The meta queries (stored in `self.transformer.meta_queries`) are inserted after all valid tokens
+ (as indicated by attention_mask) but before padding tokens.
+ The attention mask is updated with ones for the meta queries.
+
+ Args:
+ inputs_embeds (torch.Tensor): Input embeddings of shape (B, L, d_model).
+ attention_mask (torch.Tensor): Attention mask of shape (B, L), with 1 for valid tokens and 0 for padding.
+ image_mask (Optional[torch.Tensor]): Image mask of shape (B, L).
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ - Updated input embeddings of shape (B, L + N_queries, d_model)
+ - Updated attention mask of shape (B, L + N_queries)
+ - Updated image mask of shape (B, L + N_queries)
+ """
+ batch_size, _, _ = inputs_embeds.size()
+ num_meta_queries = self.transformer.meta_queries.size(0) # Number of meta query tokens
+
+ # Lists to collect updated tensors for each sample.
+ updated_embeds_list = []
+ updated_attention_list = []
+ updated_image_mask_list = []
+
+ for batch_idx in range(batch_size):
+ # Compute the number of valid tokens (non-padding) using the attention mask.
+ valid_length = int(attention_mask[batch_idx].sum().item())
+
+ # Insert meta queries after valid tokens and before any padding tokens.
+ valid_embeds = inputs_embeds[batch_idx, :valid_length, :] # (valid_length, d_model)
+ pad_embeds = inputs_embeds[batch_idx, valid_length:, :] # (pad_length, d_model)
+ updated_embeds = torch.cat([valid_embeds, self.transformer.meta_queries, pad_embeds], dim=0)
+ updated_embeds_list.append(updated_embeds)
+
+ # Update the attention mask:
+ valid_mask = attention_mask[batch_idx, :valid_length] # (valid_length,) all ones
+ pad_mask = attention_mask[batch_idx, valid_length:] # (pad_length,) all zeros
+ meta_mask = torch.ones(num_meta_queries, device=attention_mask.device, dtype=attention_mask.dtype)
+ updated_attention = torch.cat([valid_mask, meta_mask, pad_mask], dim=0)
+ updated_attention_list.append(updated_attention)
+
+ if image_mask is not None:
+ valid_image_mask = image_mask[batch_idx, :valid_length] # (valid_length,) all ones
+ pad_image_mask = image_mask[batch_idx, valid_length:] # (pad_length,) all zeros
+ meta_image_mask = torch.zeros(num_meta_queries, device=image_mask.device, dtype=image_mask.dtype)
+ updated_image_mask = torch.cat([valid_image_mask, meta_image_mask, pad_image_mask], dim=0)
+ updated_image_mask_list.append(updated_image_mask)
+
+ # Stack lists to form batch tensors.
+ updated_embeds = torch.stack(updated_embeds_list, dim=0) # (B, L + N_queries, d_model)
+ updated_attention_mask = torch.stack(updated_attention_list, dim=0) # (B, L + N_queries)
+
+ if updated_image_mask_list:
+ updated_image_mask = torch.stack(updated_image_mask_list, dim=0) # (B, L + N_queries)
+ else:
+ updated_image_mask = None # type: ignore
+
+ return updated_embeds, updated_attention_mask, updated_image_mask
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor,
+ mm_token_type_ids: torch.IntTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.LongTensor:
+ """Get the rope index.
+
+ Args:
+ input_ids (torch.LongTensor): The input IDs of the model.
+ mm_token_type_ids (torch.IntTensor): Token modality ids for text, image, and video tokens.
+ image_grid_thw (torch.LongTensor | None): The grid size of the image.
+ video_grid_thw (torch.LongTensor | None): The grid size of the video.
+ attention_mask (torch.Tensor | None): The attention mask of the model.
+
+ Returns:
+ torch.LongTensor: The position IDs.
+ """
+ attention_mask_tensor = (
+ attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"] # type: ignore
+ )
+ if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
+ attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
+ # Only apply conversion for floating point tensors (inverted masks)
+ if attention_mask_tensor.dtype.is_floating_point:
+ attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min
+ attention_mask_tensor = (1.0 - attention_mask_tensor).int()
+
+ # Calculate RoPE index once per generation.
+ position_ids, rope_deltas = self.text_encoder.model.get_rope_index(
+ input_ids,
+ mm_token_type_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask=attention_mask_tensor,
+ )
+ self.text_encoder.model.rope_deltas = rope_deltas
+ return position_ids
+
+ def prepare_initial_input_embeddings(
+ self,
+ input_ids: torch.LongTensor,
+ attention_mask: torch.Tensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ **extra_model_inputs: torch.Tensor,
+ ) -> dict[str, torch.Tensor]:
+ """Prepare initial input embeddings for the model.
+
+ Args:
+ input_ids (torch.LongTensor): The input IDs of the model.
+ attention_mask (torch.Tensor): The attention mask of the model.
+ image_grid_thw (torch.LongTensor | None): The grid size of the image.
+ pixel_values (torch.Tensor | None): The pixel values of the image.
+ extra_model_inputs (dict[str, torch.Tensor]): Additional processor outputs that should be forwarded to the
+ language model unchanged, such as multimodal token type ids.
+
+ Returns:
+ dict[str, torch.Tensor]: The model inputs:
+ - inputs_embeds: The initial input embeddings.
+ - deepstack_visual_embeds: The deepstack visual embeds or None if no image is provided.
+ - position_ids: The position ids.
+ - attention_mask: The attention mask.
+ - visual_pos_masks: The visual pos masks or None if no image is provided.
+ """
+ text_encoder_device = self.text_encoder.get_input_embeddings().weight.device
+ input_ids = input_ids.to(text_encoder_device)
+ attention_mask = attention_mask.to(text_encoder_device)
+ image_grid_thw = image_grid_thw.to(text_encoder_device) if image_grid_thw is not None else None
+ pixel_values = pixel_values.to(text_encoder_device) if pixel_values is not None else None
+ extra_model_inputs = {
+ key: value.to(text_encoder_device) if hasattr(value, "to") else value
+ for key, value in extra_model_inputs.items()
+ }
+
+ inputs_embeds = self.text_encoder.get_input_embeddings()(input_ids)
+ attention_mask = attention_mask.to(inputs_embeds.device)
+ attention_mask = attention_mask.bool()
+
+ if pixel_values is not None:
+ image_features = self.text_encoder.get_image_features(pixel_values, image_grid_thw)
+ if hasattr(image_features, "pooler_output"):
+ image_embeds = image_features.pooler_output
+ deepstack_visual_embeds = image_features.deepstack_features
+ else:
+ image_embeds, deepstack_visual_embeds = image_features
+ image_embeds = torch.cat(image_embeds, dim=0)
+ image_mask, _ = self.text_encoder.model.get_placeholder_mask(
+ input_ids,
+ inputs_embeds=inputs_embeds,
+ image_features=image_embeds,
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ image_mask = image_mask[..., 0]
+ visual_pos_masks = image_mask
+ else:
+ visual_pos_masks = None
+ deepstack_visual_embeds = None
+
+ # Add meta queries to the input embeddings
+ inputs_embeds, attention_mask, visual_pos_masks = self.prepare_inputs_for_meta_queries( # type: ignore
+ inputs_embeds,
+ attention_mask,
+ image_mask=visual_pos_masks,
+ )
+
+ if "mm_token_type_ids" in extra_model_inputs:
+ mm_token_type_ids = extra_model_inputs["mm_token_type_ids"]
+ updated_mm_token_type_ids = []
+ for batch_idx in range(mm_token_type_ids.shape[0]):
+ valid_length = int(attention_mask[batch_idx].sum().item()) - self.transformer.meta_queries.size(0)
+ valid_token_types = mm_token_type_ids[batch_idx, :valid_length]
+ pad_token_types = mm_token_type_ids[batch_idx, valid_length:]
+ meta_token_types = torch.zeros(
+ self.transformer.meta_queries.size(0),
+ device=mm_token_type_ids.device,
+ dtype=mm_token_type_ids.dtype,
+ )
+ updated_mm_token_type_ids.append(torch.cat([valid_token_types, meta_token_types, pad_token_types], dim=0))
+ extra_model_inputs["mm_token_type_ids"] = torch.stack(updated_mm_token_type_ids, dim=0)
+
+ # add placeholder for meta queries, it will be used to create position ids for meta queries only
+ img_input_ids = torch.full(
+ (input_ids.shape[0], self.transformer.meta_queries.size(0)),
+ -1,
+ device=input_ids.device,
+ dtype=torch.long,
+ )
+ input_ids = torch.cat([input_ids, img_input_ids], dim=1) # type: ignore
+
+ # computing rope index
+ position_ids = self.get_rope_index(
+ input_ids,
+ extra_model_inputs["mm_token_type_ids"],
+ image_grid_thw,
+ extra_model_inputs.get("video_grid_thw"),
+ attention_mask,
+ )
+
+ return {
+ "inputs_embeds": inputs_embeds,
+ "deepstack_visual_embeds": deepstack_visual_embeds,
+ "position_ids": position_ids,
+ "attention_mask": attention_mask,
+ "visual_pos_masks": visual_pos_masks,
+ **extra_model_inputs,
+ }
+
+ def _get_editing_embeddings(
+ self,
+ edit_prompt: list[str] | str,
+ conditioning_image: Image.Image | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Get the editing embeddings for the given prompt and image.
+
+ Args:
+ edit_prompt (list[str] | str): The editing prompt.
+ conditioning_image (Image.Image | None): The conditioning image or None if t2i generation is enabled.
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Tuple of prompt embeds and negative prompt embeds.
+ """
+ # Prepare inputs for the text encoder call.
+ prepared_inputs = self.prepare_inputs_batch(instruction=edit_prompt, image=conditioning_image)
+ model_inputs = self.prepare_initial_input_embeddings(**prepared_inputs) # type: ignore[arg-type]
+
+ # Forward the text encoder.
+ llm_model_result = self.text_encoder.model.language_model(input_ids=None, **model_inputs)
+ hidden_states = llm_model_result.last_hidden_state
+ attention_mask = model_inputs["attention_mask"]
+
+ # Get the hidden states for the meta tokens
+ meta_hidden_states = self.transformer.get_hidden_states_for_meta_tokens(hidden_states, attention_mask)
+
+ # Forward the edit heads
+ batch_size = meta_hidden_states.shape[0]
+ prompt_embeds, negative_prompt_embeds = self.transformer.forward_edit_heads(batch_size, meta_hidden_states)
+ target_device = self._execution_device
+ return prompt_embeds.to(target_device), negative_prompt_embeds.to(target_device)
+
+ def encode_prompt( # type: ignore[override]
+ self,
+ batch_size: int,
+ prompt: str | list[str] | None = None,
+ conditioning_image: Image.Image | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ num_images_per_prompt: int = 1,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Encode the prompt into text encoder hidden states.
+
+ Args:
+ batch_size (int): The batch size.
+ conditioning_image: Image.Image: The conditioning image to pass in prompt encoder.
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ Number of images that should be generated per prompt. torch device to place the resulting embeddings on
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Tuple of prompt embeds and negative prompt embeds.
+ """
+ if prompt_embeds is None:
+ prompt_embeds, negative_prompt_embeds = self._get_editing_embeddings(
+ edit_prompt=prompt, # type: ignore
+ conditioning_image=conditioning_image,
+ )
+
+ seq_len = prompt_embeds.size(1)
+ 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)
+
+ if self.do_classifier_free_guidance:
+ if negative_prompt_embeds is None:
+ msg = "Negative prompt embeds are required when classifier free guidance is enabled."
+ raise ValueError(msg)
+ seq_len = negative_prompt_embeds.size(1)
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+ if self.do_image_guidance:
+ prompt_embeds = torch.cat( # type: ignore
+ [prompt_embeds, negative_prompt_embeds, negative_prompt_embeds],
+ dim=0,
+ )
+ else:
+ prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0)
+
+ return prompt_embeds, negative_prompt_embeds # type: ignore
+
+ def prepare_guidance_inputs(
+ self,
+ image_latents: torch.Tensor,
+ prompt_embeds: torch.Tensor,
+ negative_prompt_embeds: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Prepare the guidance inputs for the denoising process.
+
+ Args:
+ image_latents (torch.Tensor): The image latents to be used for the denoising process.
+ prompt_embeds (torch.Tensor): The prompt embeds from edit head.
+ negative_prompt_embeds (torch.Tensor): The negative prompt embeds from edit head.
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Tuple of image latents and prompt embeds for the denoising process.
+ """
+ if self.do_classifier_free_guidance:
+ uncond_image_latents = torch.zeros_like(image_latents)
+ image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0)
+
+ prompt_embeds = torch.cat(
+ [
+ prompt_embeds,
+ negative_prompt_embeds,
+ negative_prompt_embeds,
+ ],
+ dim=0,
+ )
+ return image_latents, prompt_embeds
+
+ def prepare_image_latents(
+ self,
+ image: torch.Tensor,
+ batch_size: int,
+ num_images_per_prompt: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ ) -> torch.Tensor:
+ """Prepare the coonditioning image latents for the denoising process.
+
+ Args:
+ image (torch.Tensor): The image to process.
+ batch_size (int): The batch size.
+ num_images_per_prompt (int): The number of images to generate per prompt.
+ dtype (torch.dtype): The dtype of the image.
+ device (torch.device): The device to use.
+ """
+ image = image.to(device=device, dtype=dtype)
+ batch_size = batch_size * num_images_per_prompt
+
+ image_latents = self.vae.encode(image).latent
+ image_latents = image_latents * self.vae.config.scaling_factor
+
+ image_latents = torch.cat([image_latents] * batch_size, dim=0)
+
+ if self.do_classifier_free_guidance:
+ if self.do_image_guidance:
+ uncond_image_latents = torch.zeros_like(image_latents)
+ image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0)
+ else:
+ image_latents = torch.cat([image_latents, image_latents], dim=0)
+ return image_latents
+
+ def _run_denoising_loop(
+ self,
+ latents: torch.Tensor,
+ latent_channels: int,
+ input_image_latents: torch.Tensor | None,
+ timesteps: torch.Tensor,
+ num_inference_steps: int,
+ num_warmup_steps: int,
+ prompt_embeds: torch.Tensor,
+ guidance_scale: float,
+ image_guidance_scale: float,
+ extra_step_kwargs: dict,
+ *,
+ is_t2i: bool = False,
+ ) -> torch.Tensor:
+ """Run the denoising loop over the given timesteps.
+
+ Args:
+ latents (torch.Tensor): Latents to denoise.
+ latent_channels (int): The number of latent channels.
+ input_image_latents (torch.Tensor | None): The input image latents or None if t2i generation is enabled.
+ timesteps (torch.Tensor): The timesteps for the scheduler.
+ num_inference_steps (int): The number of total inference steps.
+ num_warmup_steps (int): The number of warmup steps for the scheduler.
+ prompt_embeds (torch.Tensor): The prompt embeds from edit head.
+ guidance_scale (float): The guidance scale for the text prompt.
+ image_guidance_scale (float): The guidance scale for the image prompt.
+ extra_step_kwargs (dict): The extra step kwargs for the scheduler.
+ is_t2i (bool): Whether the t2i generation is enabled.
+
+ Returns:
+ torch.Tensor: The final latents after the denoising loop.
+ """
+ transformer_dtype = self.transformer.dtype
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ if self.do_classifier_free_guidance:
+ latent_model_input = torch.cat([latents] * (3 if self.do_image_guidance else 2))
+ else:
+ latent_model_input = latents
+
+ # concat noised and input image latents.
+ if input_image_latents is None: # In case of t2i generation, we don't have input image latents
+ scaled_latent_model_input = torch.cat([latent_model_input, latent_model_input], dim=1)
+ else:
+ scaled_latent_model_input = torch.cat([latent_model_input, input_image_latents], dim=1)
+
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latent_model_input.shape[0])
+ timestep = timestep * self.transformer.config.timestep_scale
+
+ # predict noise model_output
+ noise_pred = self.transformer(
+ hidden_states=scaled_latent_model_input.to(dtype=transformer_dtype),
+ encoder_hidden_states=prompt_embeds,
+ timestep=timestep,
+ return_dict=False,
+ attention_kwargs=self.attention_kwargs,
+ t2i_samples=[is_t2i] * scaled_latent_model_input.shape[0],
+ )[0]
+ noise_pred = noise_pred.float()
+
+ # perform guidance
+ if self.do_classifier_free_guidance:
+ if self.do_image_guidance:
+ noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3)
+ noise_pred = (
+ noise_pred_uncond
+ + guidance_scale * (noise_pred_text - noise_pred_image)
+ + image_guidance_scale * (noise_pred_image - noise_pred_uncond)
+ )
+ else:
+ noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+ # learned sigma
+ if self.transformer.config.out_channels // 2 == latent_channels:
+ noise_pred = noise_pred.chunk(2, dim=1)[0]
+
+ # compute previous image: x_t -> x_t-1
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ return latents
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ conditioning_image: Image.Image | None = None,
+ prompt: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ latents: torch.FloatTensor | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ eta: float = 0.0,
+ num_inference_steps: int = 20,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float = 4.5,
+ image_guidance_scale: float = 1.2,
+ num_images_per_prompt: int = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str = "pil",
+ attention_kwargs: dict[str, Any] | None = None,
+ *,
+ return_dict: bool = True,
+ use_resolution_binning: bool = True,
+ **_: Any,
+ ) -> SanaPipelineOutput | tuple:
+ """Function invoked when calling the pipeline for generation.
+
+ Args:
+ conditioning_image (Image.Image | None): The input conditioning image or None if t2i generation is enabled.
+ prompt (str | list[str] | None): The editing prompt.
+ prompt_embeds (torch.FloatTensor): The prompt embeds from edit head.
+ negative_prompt_embeds (torch.FloatTensor): The negative prompt embeds from edit head.
+ latents (Optional[torch.FloatTensor]): The latents to use for the denoising process.
+ height (Optional[int]): The height of the conditioning_image.
+ width (Optional[int]): The width of the conditioning_image.
+ num_inference_steps (int): The number of inference steps.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+ passed will be used. Must be in descending order.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+ will be used.
+ eta (`float`, *optional*, defaults to 0.0):
+ Corresponds to parameter eta (η) in the DDIM paper: huggingface.co/papers/2010.02502. Only
+ applies to [`schedulers.DDIMScheduler`], will be ignored for others.
+ guidance_scale (float): The guidance scale for the text prompt.
+ image_guidance_scale (float): The guidance scale for the conditioning_image.
+ num_images_per_prompt (int): The number of images to generate per prompt.
+ generator (Optional[Union[torch.Generator, List[torch.Generator]]]): The generator.
+ output_type (str): The output type.
+ attention_kwargs (dict[str, Any]): The attention kwargs.
+ return_dict (bool): Whether to return a return_dict.
+ use_resolution_binning (bool): Whether to use resolution binning.
+ _: Additional keyword arguments.
+
+ Returns:
+ SanaPipelineOutput: The output of the pipeline.
+ """
+ # check if we need to fallback to t2i generation
+ is_t2i = False
+ if conditioning_image is None:
+ is_t2i = True
+
+ # 0. Set pipeline attributes.
+ device = self._execution_device
+ dtype = self.dtype # type: ignore
+ self._guidance_scale = guidance_scale
+ self._image_guidance_scale = 0 if is_t2i else image_guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._interrupt = False
+
+ if (height is None or width is None) and conditioning_image is None:
+ msg = "Either height and width or conditioning_image must be provided."
+ raise ValueError(msg)
+
+ # 1. Check inputs. Raise error if not correct
+ if use_resolution_binning:
+ if height is None or width is None:
+ height, width = conditioning_image.height, conditioning_image.width # type: ignore[union-attr]
+ orig_height, orig_width = height, width
+ height, width = self.get_closest_size(height, width)
+
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ )
+
+ # Identify batch size.
+ 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] # type: ignore
+
+ # 1. Encode input prompt
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt( # type: ignore[assignment]
+ batch_size=batch_size,
+ conditioning_image=conditioning_image,
+ prompt=prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+
+ # 2. Prepare conditioning image latents
+ if is_t2i:
+ image_latents = None
+ else:
+ processed_image = self.image_processor.preprocess(conditioning_image, height=height, width=width)
+ processed_image = processed_image.to(device).to(dtype)
+ image_latents = self.prepare_image_latents(
+ processed_image,
+ batch_size,
+ num_images_per_prompt,
+ self.dtype,
+ device,
+ )
+
+ # 3. Prepare latents
+ latent_channels = self.vae.config.latent_channels
+ latents = self.prepare_latents( # type: ignore
+ batch_size * num_images_per_prompt,
+ latent_channels,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents,
+ )
+
+ num_channels_latents_check = (
+ latent_channels * 2 if self.transformer.input_condition_type == "channel_cat" else latent_channels
+ )
+ if num_channels_latents_check != self.transformer.config.in_channels:
+ msg = f"The config of `pipeline.transformer` expects {self.transformer.config.in_channels} channels,"
+ f"but received {num_channels_latents_check}."
+ raise ValueError(msg)
+
+ # 4. Prepare timesteps
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ device,
+ timesteps,
+ sigmas,
+ )
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) # type: ignore[arg-type]
+ self._num_timesteps = len(timesteps) # type: ignore[arg-type]
+
+ # 5. Prepare extra step kwargs.
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ # 6. Denoising loop
+ latents = self._run_denoising_loop( # type: ignore
+ latents=latents, # type: ignore[arg-type]
+ latent_channels=latent_channels,
+ input_image_latents=image_latents,
+ timesteps=timesteps, # type: ignore[arg-type]
+ num_inference_steps=num_inference_steps,
+ num_warmup_steps=num_warmup_steps,
+ prompt_embeds=prompt_embeds, # type: ignore[arg-type]
+ guidance_scale=guidance_scale,
+ image_guidance_scale=image_guidance_scale,
+ extra_step_kwargs=extra_step_kwargs,
+ is_t2i=is_t2i,
+ )
+
+ if output_type == "latent":
+ image = latents # type: ignore[assignment]
+ else:
+ latents = latents.to(self.dtype) # type: ignore
+ image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+
+ if use_resolution_binning:
+ image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
+
+ if output_type != "latent":
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (image,)
+
+ return SanaPipelineOutput(images=image) # type: ignore[arg-type]
+
+
+class VIBESanaImagePipeline(VIBESanaEditingPipeline):
+ def __call__(
+ self,
+ image: Image.Image | None = None,
+ prompt: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ latents: torch.FloatTensor | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ eta: float = 0.0,
+ num_inference_steps: int = 20,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float = 4.5,
+ image_guidance_scale: float = 1.2,
+ num_images_per_prompt: int = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str = "pil",
+ attention_kwargs: dict[str, Any] | None = None,
+ *,
+ return_dict: bool = True,
+ use_resolution_binning: bool = True,
+ **_: Any,
+ ) -> SanaPipelineOutput | tuple:
+ return super().__call__(
+ conditioning_image=image[0],
+ prompt=prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ latents=latents,
+ height=height,
+ width=width,
+ eta=eta,
+ num_inference_steps=num_inference_steps,
+ timesteps=timesteps,
+ sigmas=sigmas,
+ guidance_scale=guidance_scale,
+ image_guidance_scale=image_guidance_scale,
+ num_images_per_prompt=num_images_per_prompt,
+ generator=generator,
+ output_type=output_type,
+ attention_kwargs=attention_kwargs,
+ return_dict=return_dict,
+ use_resolution_binning=use_resolution_binning,
+ )
diff --git a/pipelines/xomni/modeling_siglip_flux.py b/pipelines/xomni/modeling_siglip_flux.py
index 18b7b3463..6f1754b3c 100644
--- a/pipelines/xomni/modeling_siglip_flux.py
+++ b/pipelines/xomni/modeling_siglip_flux.py
@@ -291,8 +291,8 @@ def teacache_forward(
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
)
- batch_size, seq_len, channels = hidden_states.shape
- device, dtype = hidden_states.device, hidden_states.dtype
+ _batch_size, _seq_len, _channels = hidden_states.shape
+ _device, _dtype = hidden_states.device, hidden_states.dtype
hidden_states = self.x_embedder(hidden_states)
timestep = timestep.to(hidden_states.dtype) * 1000
@@ -332,7 +332,7 @@ def teacache_forward(
if self.enable_teacache:
inp = hidden_states.clone()
temb_ = temb.clone()
- modulated_inp, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.transformer_blocks[0].norm1(inp, emb=temb_)
+ modulated_inp, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = self.transformer_blocks[0].norm1(inp, emb=temb_)
if self.cnt == 0 or self.cnt == self.num_steps - 1:
should_calc = True
self.accumulated_rel_l1_distance = 0
@@ -502,9 +502,9 @@ class FluxPipelineWithSigLIP(FluxPipeline):
def __call__(
self,
siglip_tensor: torch.Tensor,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ negative_prompt: Union[str, List[str]] | None = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
true_cfg_scale: float = 1.0,
true_cfg_scale_2: float = 1.0,
diff --git a/pipelines/xomni/modeling_siglip_tokenizer.py b/pipelines/xomni/modeling_siglip_tokenizer.py
index 5e99d4ff8..16f79d865 100644
--- a/pipelines/xomni/modeling_siglip_tokenizer.py
+++ b/pipelines/xomni/modeling_siglip_tokenizer.py
@@ -78,8 +78,8 @@ class IBQ(nn.Module):
def forward(self, z, temp=None, rescale_logits=False, return_logits=False, **kwargs):
assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel"
- assert rescale_logits == False, "Only for interface compatible with Gumbel"
- assert return_logits == False, "Only for interface compatible with Gumbel"
+ assert not rescale_logits, "Only for interface compatible with Gumbel"
+ assert not return_logits, "Only for interface compatible with Gumbel"
# reshape z -> (batch, height, width, channel) and flatten
z = rearrange(z, 'b c h w -> b h w c').contiguous()
assert z.shape[-1] == self.e_dim
diff --git a/pipelines/xomni/modeling_vit.py b/pipelines/xomni/modeling_vit.py
index 150571c1d..dfd8a1da1 100644
--- a/pipelines/xomni/modeling_vit.py
+++ b/pipelines/xomni/modeling_vit.py
@@ -39,7 +39,7 @@ def _no_grad_trunc_normal_(tensor, mean, std, a, b):
# Values are generated by using a truncated uniform distribution and
# then using the inverse CDF for the normal distribution.
# Get upper and lower cdf values
- l = norm_cdf((a - mean) / std) # noqa: E741
+ l = norm_cdf((a - mean) / std)
u = norm_cdf((b - mean) / std)
# Uniformly fill tensor with values from [l, u], then translate to
@@ -132,7 +132,7 @@ class Attention(nn.Module):
self.proj_drop = nn.Dropout(proj_drop) if proj_drop > 0.0 else nn.Identity()
def forward(self, x: torch.Tensor, cu_slens=None) -> torch.Tensor:
- B, N, C = x.shape
+ B, N, _C = x.shape
qkv = (
self.qkv(x)
.reshape(B, N, 3, self.num_heads, self.head_dim)
@@ -314,7 +314,6 @@ class VisionTransformer(nn.Module):
super().__init__()
assert global_pool in ("", "avg", "token", "map")
assert class_token or global_pool != "token"
- use_fc_norm = global_pool == "avg" if fc_norm is None else fc_norm
# norm_layer = get_norm_layer(norm_layer) or partial(nn.LayerNorm, eps=1e-6)
# act_layer = get_act_layer(act_layer) or nn.GELU
norm_layer = partial(nn.LayerNorm, eps=1e-6)
@@ -427,9 +426,7 @@ class VisionTransformer(nn.Module):
if global_pool is not None:
assert global_pool in ("", "avg", "token", "map")
if global_pool == "map" and self.attn_pool is None:
- assert (
- False
- ), "Cannot currently add attention pooling in reset_classifier()."
+ raise AssertionError("Cannot currently add attention pooling in reset_classifier().")
elif global_pool != "map " and self.attn_pool is not None:
self.attn_pool = None # remove attention pooling
self.global_pool = global_pool
@@ -539,14 +536,14 @@ class VisionTransformer(nn.Module):
x_all = []
image_sizes = []
for x in x_list:
- bs, _, h, w = x.shape
+ _bs, _, h, w = x.shape
# fix patch size=14 in datasets
pad_h = (self.patch_embed.patch_size[0] - h % self.patch_embed.patch_size[0]) % self.patch_embed.patch_size[0]
pad_w = (self.patch_embed.patch_size[1] - w % self.patch_embed.patch_size[1]) % self.patch_embed.patch_size[1]
x = F.pad(x, (0, pad_w, 0, pad_h))
- bs, _, h, w = x.shape
+ _bs, _, h, w = x.shape
h = h // self.patch_embed.patch_size[0]
w = w // self.patch_embed.patch_size[1]
@@ -566,7 +563,7 @@ class VisionTransformer(nn.Module):
cu_indices.append(cu_indices[-1] + i)
cu_slens = torch.tensor(cu_indices, dtype=torch.int32).to(x.device)
- for idx, blk in enumerate(self.blocks):
+ for _idx, blk in enumerate(self.blocks):
if self.grad_checkpointing and not torch.jit.is_scripting():
x = checkpoint(blk, x, cu_slens, use_reentrant=True)
else:
@@ -575,7 +572,7 @@ class VisionTransformer(nn.Module):
return feats, image_sizes
def forward_features(self, x: torch.Tensor) -> torch.Tensor:
- bs, _, h, w = x.shape
+ _bs, _, h, w = x.shape
h = h // self.patch_embed.patch_size[0]
w = w // self.patch_embed.patch_size[1]
diff --git a/pyproject.toml b/pyproject.toml
index f73c5ab0e..6266072d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,36 +14,7 @@ exclude = [
".git",
".ruff_cache",
".vscode",
- "modules/cfgzero",
- "modules/flash_attn_triton_amd",
- "modules/hidiffusion",
- "modules/intel/ipex",
- "modules/pag",
- "modules/schedulers",
- "modules/teacache",
- "modules/seedvr",
- "modules/sharpfin",
- "modules/control/proc",
- "modules/control/units",
- "modules/control/units/xs_pipe.py",
- "modules/postprocess/aurasr_arch.py",
- "pipelines/meissonic",
- "pipelines/omnigen2",
- "pipelines/hdm",
- "pipelines/segmoe",
- "pipelines/xomni",
- "pipelines/chrono",
- "scripts/lbm",
- "scripts/daam",
- "scripts/xadapter",
- "scripts/pulid",
- "scripts/instantir",
- "scripts/freescale",
- "scripts/consistory",
- "extensions-builtin/Lora",
- "extensions-builtin/sd-extension-chainner/nodes",
- "extensions-builtin/sd-webui-agent-scheduler",
- "extensions-builtin/sdnext-modernui/node_modules",
+ "modules/seedvr/src",
]
[tool.ruff.lint]
@@ -74,34 +45,37 @@ ignore = [
"ASYNC240", # Async functions should not use os.path methods
"B006", # Do not use mutable data structures for argument defaults
"B008", # Do not perform function call in argument defaults
- "B905", # Strict zip() usage
+ "B019", # Do not use `lru_cache`
+ "B026", # Star arg unpacking
+ "B905", # Strict zip usage
+ "C401", # Unnecessary `generator` comprehension
"C408", # Unnecessary `dict` call
- "C420", # Unnecessary dict comprehension for iterable; use `dict.fromkeys` instead
+ "C417", # Unnecessary `map` comprehension
+ "C420", # Unnecessary `dict` comprehension
"E402", # Module level import not at top of file
"E501", # Line too long
- "E721", # Do not compare types, use `isinstance()`
- "E731", # Do not assign a `lambda` expression, use a `def`
+ "E721", # Do not compare types
+ "E731", # Do not assign a `lambda` expression
"E741", # Ambiguous variable name
- "EXE001", # file with shebang is not marked executable
- "F401", # Imported by unused
- "I001", # Import block is un-sorted or un-formatted
+ "F401", # Import unused
+ "F403", # Import star
+ "I001", # Import block is un-sorted
"NPY002", # replace legacy random
"RUF005", # Consider iterable unpacking
- "RUF008", # Do not use mutable default values for dataclass
"RUF010", # Use explicit conversion flag
"RUF012", # Mutable class attributes
- "RUF015", # Prefer `next(...)` over single element slice
+ "RUF015", # Prefer `next`
"RUF022", # All is not sorted
"RUF046", # Value being cast to `int` is already an integer
"RUF051", # Prefer pop over del
- "RUF059", # Unpacked variables are not used
]
fixable = ["ALL"]
unfixable = []
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
[tool.ruff.lint.per-file-ignores]
-"modules/caption/joytag.py" = ["RUF013"] # Per header comment: "Do not modify directly — sync from upstream"
+"modules/seedvr/**/*.py" = ["TID252"] # Vendored code: relative imports required for package structure
+"pipelines/omnigen2/**/*.py" = ["TID252"] # Vendored code: relative imports required for package structure
[tool.ruff.format]
quote-style = "double"
@@ -130,10 +104,12 @@ main.ignore-paths=[
".git",
".ruff_cache",
".vscode",
+ ".*/node_modules/.*",
+ "modules/control/units/xs_pipe.py",
+ "modules/control/proc",
+ "modules/schedulers/scheduler_*.py",
"modules/apg",
"modules/cfgzero",
- "modules/control/proc",
- "modules/control/units",
"modules/dml",
"modules/face",
"modules/flash_attn_triton_amd",
@@ -151,7 +127,6 @@ main.ignore-paths=[
"modules/seedvr",
"modules/sharpfin",
"modules/rife",
- "modules/schedulers",
"modules/taesd",
"modules/teacache",
"modules/todo",
@@ -166,6 +141,10 @@ main.ignore-paths=[
"pipelines/segmoe",
"pipelines/xomni",
"pipelines/chrono",
+ "pipelines/step1x",
+ "pipelines/vibe",
+ "pipelines/ultraflux",
+ "pipelines/lumina_dimmo",
"scripts/consistory",
"scripts/ctrlx",
"scripts/daam",
@@ -180,11 +159,7 @@ main.ignore-paths=[
"scripts/differential_diffusion.py",
"scripts/pulid",
"scripts/xadapter",
- "repositories",
"extensions-builtin/sd-extension-chainner/nodes",
- "extensions-builtin/sd-webui-agent-scheduler",
- "extensions-builtin/sdnext-modernui/node_modules",
- "extensions-builtin/sdnext-kanvas/node_modules",
]
main.ignore-patterns=[
".*test*.py$",
@@ -392,8 +367,6 @@ exclude = [
"venv",
]
extraPaths = [
- "scripts",
- "pipelines",
]
reportMissingImports = "none"
reportInvalidTypeForm = "none"
@@ -414,6 +387,36 @@ include = [
exclude = [
"venv/",
"*.git/",
+ "scripts/pulid/*",
+ "scripts/pixelsmith/*",
+ "scripts/mod/*",
+ "scripts/layerdiffuse/*",
+ "scripts/lbm/*",
+ "scripts/daam/*",
+ "scripts/infiniteyou/*",
+ "scripts/ctrlx/*",
+ "scripts/consistory/*",
+ "scripts/freescale/*",
+ "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/bria",
+ "pipelines/ernie",
+ "pipelines/flex2",
+ "pipelines/anima",
+ "pipelines/meissonic",
+ "extensions-builtin/sd-extension-chainner/nodes",
]
[tool.ty.rules]
@@ -428,3 +431,9 @@ unsupported-base = "ignore"
possibly-missing-attribute = "ignore"
invalid-parameter-default = "ignore"
call-non-callable = "ignore"
+not-subscriptable = "ignore"
+possibly-missing-submodule = "ignore"
+missing-argument = "ignore"
+unknown-argument = "ignore"
+invalid-attribute-access = "ignore"
+call-top-callable = "ignore"
diff --git a/requirements.txt b/requirements.txt
index dbdf506a1..f1391e86a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -23,14 +23,14 @@ ftfy
# versioned
fastapi==0.124.4
rich==14.1.0
-safetensors==0.7.0
+safetensors==0.8.0rc0
peft==0.19.1
httpx==0.28.1
requests==2.32.3
tqdm==4.67.3
accelerate==1.13.0
einops==0.8.1
-huggingface_hub==1.12.0
+huggingface_hub==1.14.0
numpy==2.1.2
pandas==2.3.1
protobuf==6.33.5
diff --git a/scripts/autocomplete.py b/scripts/autocomplete.py
index 0665ace6a..664619afd 100644
--- a/scripts/autocomplete.py
+++ b/scripts/autocomplete.py
@@ -27,7 +27,7 @@ def get_config_json():
"autocomplete_active": bool(shared.opts.data.get('autocomplete_active', False)),
"autocomplete_enabled": enabled,
"autocomplete_min_chars": shared.opts.data.get('autocomplete_min_chars', 3),
- "autocomplete_replace_underscores": shared.opts.data.get('autocomplete_replace_underscores', True),
+ "autocomplete_keep_underscores": shared.opts.data.get('autocomplete_keep_underscores', False),
"autocomplete_append_comma": shared.opts.data.get('autocomplete_append_comma', True),
"autocomplete_at_prefix_artist": shared.opts.data.get('autocomplete_at_prefix_artist', False),
"autocomplete_translations": bool(shared.opts.data.get('autocomplete_translations', False)),
@@ -52,8 +52,8 @@ def on_min_chars_change(value):
return get_config_json()
-def on_replace_underscores_change(value):
- shared.opts.data['autocomplete_replace_underscores'] = bool(value)
+def on_keep_underscores_change(value):
+ shared.opts.data['autocomplete_keep_underscores'] = bool(value)
shared.opts.save(silent=True)
return get_config_json()
@@ -180,10 +180,10 @@ class AutocompleteScript(scripts_manager.Script):
refresh_btn = ToolButton(value=symbols.refresh, elem_id=self.elem_id("refresh"))
update_btn = ToolButton(value=symbols.save, elem_id=self.elem_id("update"))
with gr.Row():
- replace_underscores = gr.Checkbox(
- label="Replace underscores",
- value=shared.opts.data.get('autocomplete_replace_underscores', True),
- elem_id=self.elem_id("replace_underscores"),
+ keep_underscores = gr.Checkbox(
+ label="Keep underscores",
+ value=shared.opts.data.get('autocomplete_keep_underscores', False),
+ elem_id=self.elem_id("keep_underscores"),
)
append_comma = gr.Checkbox(
label="Comma separator",
@@ -217,14 +217,14 @@ class AutocompleteScript(scripts_manager.Script):
active_cb.change(fn=on_active_change, inputs=[active_cb], outputs=[config_json, status])
enabled_dd.change(fn=on_enabled_change, inputs=[enabled_dd], outputs=[config_json, status])
min_chars.change(fn=on_min_chars_change, inputs=[min_chars], outputs=[config_json])
- replace_underscores.change(fn=on_replace_underscores_change, inputs=[replace_underscores], outputs=[config_json])
+ keep_underscores.change(fn=on_keep_underscores_change, inputs=[keep_underscores], outputs=[config_json])
append_comma.change(fn=on_append_comma_change, inputs=[append_comma], outputs=[config_json])
at_prefix_artist.change(fn=on_at_prefix_artist_change, inputs=[at_prefix_artist], outputs=[config_json])
translations_cb.change(fn=on_translations_change, inputs=[translations_cb], outputs=[config_json])
refresh_btn.click(fn=on_refresh, inputs=[], outputs=[enabled_dd, status])
update_btn.click(fn=on_update, inputs=[enabled_dd], outputs=[status])
- for comp in [enabled_dd, min_chars, replace_underscores, append_comma, at_prefix_artist, translations_cb, config_json, status]:
+ for comp in [enabled_dd, min_chars, keep_underscores, append_comma, at_prefix_artist, translations_cb, config_json, status]:
comp.do_not_save_to_config = True
- return [active_cb, enabled_dd, min_chars, replace_underscores, append_comma, at_prefix_artist, translations_cb, config_json]
+ return [active_cb, enabled_dd, min_chars, keep_underscores, append_comma, at_prefix_artist, translations_cb, config_json]
diff --git a/scripts/consistory/attention_processor.py b/scripts/consistory/attention_processor.py
index 04985ee4f..1aafb3f05 100644
--- a/scripts/consistory/attention_processor.py
+++ b/scripts/consistory/attention_processor.py
@@ -117,7 +117,7 @@ class ConsistoryExtendedAttnXFormersAttnProcessor:
is_cross = encoder_hidden_states is not None
perform_extend_attn = perform_extend_attn and (not is_cross) and \
- any([self.attnstore.curr_iter >= x[0] and self.attnstore.curr_iter <= x[1] for x in self.t_range]) and \
+ any(self.attnstore.curr_iter >= x[0] and self.attnstore.curr_iter <= x[1] for x in self.t_range) and \
self.curr_unet_part in self.extend_kv_unet_parts
batch_size, key_tokens, _ = (
@@ -254,19 +254,6 @@ class ConsistoryExtendedAttnXFormersAttnProcessor:
def register_extended_self_attn(unet, attnstore, extended_attn_kwargs):
- DICT_PLACE_TO_RES = {'down_0': 64, 'down_1': 64, 'down_2': 64, 'down_3': 64, 'down_4': 64, 'down_5': 64, 'down_6': 64, 'down_7': 64,
- 'down_8': 32, 'down_9': 32, 'down_10': 32, 'down_11': 32, 'down_12': 32, 'down_13': 32, 'down_14': 32, 'down_15': 32,
- 'down_16': 32, 'down_17': 32, 'down_18': 32, 'down_19': 32, 'down_20': 32, 'down_21': 32, 'down_22': 32, 'down_23': 32,
- 'down_24': 32, 'down_25': 32, 'down_26': 32, 'down_27': 32, 'down_28': 32, 'down_29': 32, 'down_30': 32, 'down_31': 32,
- 'down_32': 32, 'down_33': 32, 'down_34': 32, 'down_35': 32, 'down_36': 32, 'down_37': 32, 'down_38': 32, 'down_39': 32,
- 'down_40': 32, 'down_41': 32, 'down_42': 32, 'down_43': 32, 'down_44': 32, 'down_45': 32, 'down_46': 32, 'down_47': 32,
- 'mid_120': 32, 'mid_121': 32, 'mid_122': 32, 'mid_123': 32, 'mid_124': 32, 'mid_125': 32, 'mid_126': 32, 'mid_127': 32,
- 'mid_128': 32, 'mid_129': 32, 'mid_130': 32, 'mid_131': 32, 'mid_132': 32, 'mid_133': 32, 'mid_134': 32, 'mid_135': 32,
- 'mid_136': 32, 'mid_137': 32, 'mid_138': 32, 'mid_139': 32, 'up_49': 32, 'up_51': 32, 'up_53': 32, 'up_55': 32, 'up_57': 32,
- 'up_59': 32, 'up_61': 32, 'up_63': 32, 'up_65': 32, 'up_67': 32, 'up_69': 32, 'up_71': 32, 'up_73': 32, 'up_75': 32,
- 'up_77': 32, 'up_79': 32, 'up_81': 32, 'up_83': 32, 'up_85': 32, 'up_87': 32, 'up_89': 32, 'up_91': 32, 'up_93': 32,
- 'up_95': 32, 'up_97': 32, 'up_99': 32, 'up_101': 32, 'up_103': 32, 'up_105': 32, 'up_107': 32, 'up_109': 64, 'up_111': 64,
- 'up_113': 64, 'up_115': 64, 'up_117': 64, 'up_119': 64}
attn_procs = {}
for i, name in enumerate(unet.attn_processors.keys()):
is_self_attn = i % 2 == 0
diff --git a/scripts/consistory/consistory_pipeline.py b/scripts/consistory/consistory_pipeline.py
index 9ad065db8..a80bd5595 100644
--- a/scripts/consistory/consistory_pipeline.py
+++ b/scripts/consistory/consistory_pipeline.py
@@ -52,7 +52,7 @@ class ConsistoryExtendAttnSDXLPipeline(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
diff --git a/scripts/consistory/consistory_unet_sdxl.py b/scripts/consistory/consistory_unet_sdxl.py
index 4e6e4b335..a6694f276 100644
--- a/scripts/consistory/consistory_unet_sdxl.py
+++ b/scripts/consistory/consistory_unet_sdxl.py
@@ -788,16 +788,16 @@ class ConsistorySDXLUNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditio
b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
"""
- for i, upsample_block in enumerate(self.up_blocks):
- setattr(upsample_block, "s1", s1)
- setattr(upsample_block, "s2", s2)
- setattr(upsample_block, "b1", b1)
- setattr(upsample_block, "b2", b2)
+ for _i, upsample_block in enumerate(self.up_blocks):
+ upsample_block.s1 = s1
+ upsample_block.s2 = s2
+ upsample_block.b1 = b1
+ upsample_block.b2 = b2
def disable_freeu(self):
"""Disables the FreeU mechanism."""
freeu_keys = {"s1", "s2", "b1", "b2"}
- for i, upsample_block in enumerate(self.up_blocks):
+ for _i, upsample_block in enumerate(self.up_blocks):
for k in freeu_keys:
if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
setattr(upsample_block, k, None)
diff --git a/scripts/consistory/utils/ptp_utils.py b/scripts/consistory/utils/ptp_utils.py
index 5bd4adf65..af1e68541 100644
--- a/scripts/consistory/utils/ptp_utils.py
+++ b/scripts/consistory/utils/ptp_utils.py
@@ -99,10 +99,10 @@ class AttentionStore:
torch.cuda.empty_cache()
- def aggregate_last_steps_attention(self) -> torch.Tensor:
+ def aggregate_last_steps_attention(self):
"""Aggregates the attention across the different layers and heads at the specified resolution."""
attention_maps = torch.cat([torch.stack(x[-20:]) for x in self.step_store.values()]).mean(dim=0)
- bsz, wh, _ = attention_maps.shape
+ bsz, _wh, _ = attention_maps.shape
# Create attention maps for each concept token, for each batch item
agg_attn_maps = []
@@ -118,7 +118,7 @@ class AttentionStore:
# Upsample the attention maps to the target resolution
# and create the attention masks, unifying masks across the different concepts
for tgt_size in self.ALL_RES:
- pixels = tgt_size ** 2
+ tgt_size ** 2
tgt_agg_attn_maps = [F.interpolate(x.unsqueeze(1), size=tgt_size, mode='bilinear').squeeze(1) for x in agg_attn_maps]
attn_masks = []
@@ -189,6 +189,5 @@ class AttentionStore:
output_attn_mask[j*n_patches:(j+1)*n_patches] = attn_mask[j].unsqueeze(0) #.expand(n_patches, -1)
else:
raise NotImplementedError('mask_background_query is not supported anymore')
- output_attn_mask[0, attn_mask[i], k*n_patches:(k+1)*n_patches] = attn_mask[j].unsqueeze(0).expand(attn_mask[i].sum(), -1)
return output_attn_mask
diff --git a/scripts/consistory_ext.py b/scripts/consistory_ext.py
index e0de9f1d4..f28fc3848 100644
--- a/scripts/consistory_ext.py
+++ b/scripts/consistory_ext.py
@@ -75,7 +75,7 @@ class ConsiStoryScript(scripts_manager.Script):
shared.sd_model = sd_models.switch_pipe(cs.ConsistoryExtendAttnSDXLPipeline, shared.sd_model)
shared.sd_model.unet = cs.ConsistorySDXLUNet2DConditionModel.from_config(shared.sd_model.unet.config)
shared.sd_model.unet.load_state_dict(state_dict) # now load it into new class
- shared.sd_model.unet.to(dtype=devices.dtype)
+ shared.sd_model.unet.to(dtype=devices.dtype) # ty: ignore
state_dict = None
# sd_models.set_diffuser_options(shared.sd_model)
sd_models.move_model(shared.sd_model, devices.device)
@@ -196,7 +196,7 @@ class ConsiStoryScript(scripts_manager.Script):
log.warning(f'ConsiStory: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}')
return None
- subject, concepts, prompts, dropout, sampler, steps, same, queries, sdsa, freeu, _freeu_preset, alpha, injection = args # pylint: disable=unused-variable
+ _subject, concepts, prompts, dropout, _sampler, steps, same, queries, sdsa, _freeu, _freeu_preset, alpha, injection = args # pylint: disable=unused-variable
self.create_model() # create model if not already done
concepts, anchors, prompts, alpha, steps, seed = self.set_args(p, *args) # set arguments
diff --git a/scripts/daam/experiment.py b/scripts/daam/experiment.py
index 301efa429..e91bc1039 100644
--- a/scripts/daam/experiment.py
+++ b/scripts/daam/experiment.py
@@ -137,7 +137,7 @@ class GenerationExperiment:
(path / 'generation.pt').unlink(missing_ok=True)
- def save(self, path: str = None, heat_maps: bool = True, tokenizer: AutoTokenizer = None):
+ def save(self, path: str | None = None, heat_maps: bool = True, tokenizer: AutoTokenizer = None):
if path is None:
path = self.path
else:
@@ -166,7 +166,7 @@ class GenerationExperiment:
self.save_annotations()
- def save_annotations(self, path: Path = None):
+ def save_annotations(self, path: Path | None = None):
if path is None:
path = self.path
@@ -225,7 +225,7 @@ class GenerationExperiment:
self,
word: str,
tokenizer: PreTrainedTokenizer = None,
- crop: int = None,
+ crop: int | None = None,
output_prefix: str = '',
absolute: bool = False
) -> Path:
@@ -241,7 +241,7 @@ class GenerationExperiment:
return path
- def save_all_heat_maps(self, tokenizer: PreTrainedTokenizer = None, crop: int = None) -> Dict[str, Path]:
+ def save_all_heat_maps(self, tokenizer: PreTrainedTokenizer = None, crop: int | None = None) -> Dict[str, Path]:
path_map = {}
if tokenizer is None:
@@ -257,14 +257,14 @@ class GenerationExperiment:
return path_map
@staticmethod
- def contains_truth_mask(path: Union[str, Path], prompt_id: str = None) -> bool:
+ def contains_truth_mask(path: Union[str, Path], prompt_id: str | None = None) -> bool:
if prompt_id is None:
return any(Path(path).glob('*.gt.png'))
else:
return any((Path(path) / prompt_id).glob('*.gt.png'))
@staticmethod
- def read_seed(path: Union[str, Path], prompt_id: str = None) -> int:
+ def read_seed(path: Union[str, Path], prompt_id: str | None = None) -> int:
if prompt_id is None:
return int(Path(path).joinpath('seed.txt').read_text())
else:
@@ -279,7 +279,7 @@ class GenerationExperiment:
return (Path(path) / prompt_id / 'generation.pt').exists()
@staticmethod
- def read_prompt(path: Union[str, Path], prompt_id: str = None) -> str:
+ def read_prompt(path: Union[str, Path], prompt_id: str | None = None) -> str:
if prompt_id is None:
prompt_id = '.'
diff --git a/scripts/daam/heatmap.py b/scripts/daam/heatmap.py
index 99378f05c..0a4b07f62 100644
--- a/scripts/daam/heatmap.py
+++ b/scripts/daam/heatmap.py
@@ -64,7 +64,7 @@ def plot_overlay_heat_map(im, heat_map, word=None, out_file=None, crop=None, col
class WordHeatMap:
- def __init__(self, heatmap: torch.Tensor, word: str = None, word_idx: int = None):
+ def __init__(self, heatmap: torch.Tensor, word: str | None = None, word_idx: int | None = None):
self.word = word
self.word_idx = word_idx
self.heatmap = heatmap
@@ -128,7 +128,7 @@ class GlobalHeatMap:
self.prompt = prompt
self.compute_word_heat_map = lru_cache(maxsize=50)(self.compute_word_heat_map)
- def compute_word_heat_map(self, word: str, word_idx: int = None, offset_idx: int = 0) -> WordHeatMap:
+ def compute_word_heat_map(self, word: str, word_idx: int | None = None, offset_idx: int = 0) -> WordHeatMap:
merge_idxs, word_idx = compute_token_merge_indices(self.tokenizer, self.prompt, word, word_idx, offset_idx)
return WordHeatMap(self.heat_maps[merge_idxs].mean(0), word, word_idx)
diff --git a/scripts/daam/hook.py b/scripts/daam/hook.py
index 7c2866b63..d74b2bea3 100644
--- a/scripts/daam/hook.py
+++ b/scripts/daam/hook.py
@@ -87,7 +87,7 @@ class AggregateHooker(ObjectHooker[ModuleListType]):
class UNetCrossAttentionLocator(ModuleLocator[Attention]):
- def __init__(self, restrict: bool = None, locate_middle_block: bool = False):
+ def __init__(self, restrict: bool | None = None, locate_middle_block: bool = False):
self.restrict = restrict
self.layer_names = []
self.locate_middle_block = locate_middle_block
diff --git a/scripts/daam/trace.py b/scripts/daam/trace.py
index c748826df..1631df66c 100644
--- a/scripts/daam/trace.py
+++ b/scripts/daam/trace.py
@@ -1,12 +1,11 @@
from pathlib import Path
-from typing import List, Type, Any, Dict, Union
+from typing import List, Type, Any, Dict, Union, TYPE_CHECKING
import math
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
from diffusers.image_processor import VaeImageProcessor
from diffusers.models.attention_processor import Attention
import numpy as np
-import PIL.Image as Image
import torch
import torch.nn.functional as F
@@ -15,6 +14,9 @@ from .experiment import GenerationExperiment
from .heatmap import RawHeatMapCollection, GlobalHeatMap
from .hook import ObjectHooker, AggregateHooker, UNetCrossAttentionLocator
+if TYPE_CHECKING:
+ import PIL.Image as Image
+
__all__ = ['trace', 'DiffusionHeatMapHooker', 'GlobalHeatMap']
@@ -26,7 +28,7 @@ class DiffusionHeatMapHooker(AggregateHooker):
low_memory: bool = False,
load_heads: bool = False,
save_heads: bool = False,
- data_dir: str = None
+ data_dir: str | None = None
):
self.all_heat_maps = RawHeatMapCollection()
h = (pipeline.unet.config.sample_size * pipeline.vae_scale_factor)
@@ -117,11 +119,11 @@ class DiffusionHeatMapHooker(AggregateHooker):
try:
maps = torch.stack(all_merges, dim=0)
- except RuntimeError:
+ except RuntimeError as e:
if head_idx is not None or layer_idx is not None:
- raise RuntimeError('No heat maps found for the given parameters.')
+ raise RuntimeError('No heat maps found for the given parameters.') from e
else:
- raise RuntimeError('No heat maps found. Did you forget to call `with trace(...)` during generation?')
+ raise RuntimeError('No heat maps found. Did you forget to call `with trace(...)` during generation?') from e
maps = maps.mean(0)[:, 0]
maps = maps[:len(self.pipe.tokenizer.tokenize(prompt)) + 2] # 1 for SOS and 1 for padding
@@ -196,7 +198,7 @@ class UNetCrossAttentionHooker(ObjectHooker[Attention]):
latent_hw: int = 9216,
load_heads: bool = False,
save_heads: bool = False,
- data_dir: Union[str, Path] = None,
+ data_dir: Union[str, Path] | None = None,
):
super().__init__(module)
self.heat_maps = parent_trace.all_heat_maps
diff --git a/scripts/daam/utils.py b/scripts/daam/utils.py
index 15e2c1b5e..58402c9ed 100644
--- a/scripts/daam/utils.py
+++ b/scripts/daam/utils.py
@@ -69,7 +69,7 @@ def cache_dir() -> Path:
return Path(local, 'daam')
-def compute_token_merge_indices(tokenizer, prompt: str, word: str, word_idx: int = None, offset_idx: int = 0):
+def compute_token_merge_indices(tokenizer, prompt: str, word: str, word_idx: int | None = None, offset_idx: int = 0):
merge_idxs = []
tokens = tokenizer.tokenize(prompt.lower())
tokens = [x.replace('', '') for x in tokens] # New tokenizer uses wordpiece markers.
diff --git a/scripts/differential_diffusion.py b/scripts/differential_diffusion.py
index d0f13a078..0ca8fc463 100644
--- a/scripts/differential_diffusion.py
+++ b/scripts/differential_diffusion.py
@@ -1627,7 +1627,7 @@ class StableDiffusionDiffImg2ImgPipeline(DiffusionPipeline):
if isinstance(image[0], PIL.Image.Image):
w, h = image[0].size
- w, h = map(lambda x: x - x % 8, (w, h)) # resize to integer multiple of 8 # noqa: C417
+ w, h = map(lambda x: x - x % 8, (w, h)) # resize to integer multiple of 8
image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image]
image = np.concatenate(image, axis=0)
diff --git a/scripts/freescale/free_lunch_utils.py b/scripts/freescale/free_lunch_utils.py
index ebf165105..af64d4604 100644
--- a/scripts/freescale/free_lunch_utils.py
+++ b/scripts/freescale/free_lunch_utils.py
@@ -81,7 +81,7 @@ def register_upblock2d(model):
return forward
- for i, upsample_block in enumerate(model.unet.up_blocks):
+ for _i, upsample_block in enumerate(model.unet.up_blocks):
if isinstance_str(upsample_block, "UpBlock2D"):
upsample_block.forward = up_forward(upsample_block)
@@ -134,13 +134,13 @@ def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
return forward
- for i, upsample_block in enumerate(model.unet.up_blocks):
+ for _i, upsample_block in enumerate(model.unet.up_blocks):
if isinstance_str(upsample_block, "UpBlock2D"):
upsample_block.forward = up_forward(upsample_block)
- setattr(upsample_block, 'b1', b1)
- setattr(upsample_block, 'b2', b2)
- setattr(upsample_block, 's1', s1)
- setattr(upsample_block, 's2', s2)
+ upsample_block.b1 = b1
+ upsample_block.b2 = b2
+ upsample_block.s1 = s1
+ upsample_block.s2 = s2
def register_crossattn_upblock2d(model):
@@ -210,7 +210,7 @@ def register_crossattn_upblock2d(model):
return forward
- for i, upsample_block in enumerate(model.unet.up_blocks):
+ for _i, upsample_block in enumerate(model.unet.up_blocks):
if isinstance_str(upsample_block, "CrossAttnUpBlock2D"):
upsample_block.forward = up_forward(upsample_block)
@@ -297,10 +297,10 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
return forward
- for i, upsample_block in enumerate(model.unet.up_blocks):
+ for _i, upsample_block in enumerate(model.unet.up_blocks):
if isinstance_str(upsample_block, "CrossAttnUpBlock2D"):
upsample_block.forward = up_forward(upsample_block)
- setattr(upsample_block, 'b1', b1)
- setattr(upsample_block, 'b2', b2)
- setattr(upsample_block, 's1', s1)
- setattr(upsample_block, 's2', s2)
+ upsample_block.b1 = b1
+ upsample_block.b2 = b2
+ upsample_block.s1 = s1
+ upsample_block.s2 = s2
diff --git a/scripts/freescale/freescale_pipeline.py b/scripts/freescale/freescale_pipeline.py
index df91c014b..96748ad38 100644
--- a/scripts/freescale/freescale_pipeline.py
+++ b/scripts/freescale/freescale_pipeline.py
@@ -650,7 +650,7 @@ class StableDiffusionXLFreeScale(DiffusionPipeline, FromSingleFileMixin, LoraLoa
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
@@ -951,7 +951,7 @@ class StableDiffusionXLFreeScale(DiffusionPipeline, FromSingleFileMixin, LoraLoa
for restart_index, target_size in enumerate(target_sizes):
restart_step = restart_steps[restart_index]
- target_size_ = [target_size[0]//8, target_size[1]//8]
+ [target_size[0]//8, target_size[1]//8]
for block in self.unet.down_blocks + [self.unet.mid_block] + self.unet.up_blocks:
for module in block.modules():
@@ -1154,12 +1154,12 @@ class StableDiffusionXLFreeScale(DiffusionPipeline, FromSingleFileMixin, LoraLoa
def save_lora_weights(
self,
save_directory: Union[str, os.PathLike],
- unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
+ text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
+ text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
is_main_process: bool = True,
- weight_name: str = None,
- save_function: Callable = None,
+ weight_name: str | None = None,
+ save_function: Callable | None = None,
safe_serialization: bool = True,
):
state_dict = {}
diff --git a/scripts/freescale/freescale_pipeline_img2img.py b/scripts/freescale/freescale_pipeline_img2img.py
index 7f2964cc2..447fa054a 100644
--- a/scripts/freescale/freescale_pipeline_img2img.py
+++ b/scripts/freescale/freescale_pipeline_img2img.py
@@ -675,7 +675,7 @@ class StableDiffusionXLFreeScaleImg2Img(DiffusionPipeline, FromSingleFileMixin,
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
@@ -1210,12 +1210,12 @@ class StableDiffusionXLFreeScaleImg2Img(DiffusionPipeline, FromSingleFileMixin,
def save_lora_weights(
self,
save_directory: Union[str, os.PathLike],
- unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
+ text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
+ text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] | None = None,
is_main_process: bool = True,
- weight_name: str = None,
- save_function: Callable = None,
+ weight_name: str | None = None,
+ save_function: Callable | None = None,
safe_serialization: bool = True,
):
state_dict = {}
diff --git a/scripts/freescale/scale_attention.py b/scripts/freescale/scale_attention.py
index 9e83d5067..1ae2a2ca2 100644
--- a/scripts/freescale/scale_attention.py
+++ b/scripts/freescale/scale_attention.py
@@ -87,7 +87,7 @@ def scale_forward(
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
timestep: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
+ cross_attention_kwargs: Dict[str, Any] | None = None,
class_labels: Optional[torch.LongTensor] = None,
):
# Notice that normalization is always applied before the real computation in the following blocks.
@@ -136,7 +136,7 @@ def scale_forward(
norm_hidden_states_ = F.pad(norm_hidden_states_, (0, 0, w_jitter_range, w_jitter_range, h_jitter_range, h_jitter_range), 'constant', 0)
value = torch.zeros_like(norm_hidden_states_)
count = torch.zeros_like(norm_hidden_states_)
- for index, view in enumerate(views):
+ for _index, view in enumerate(views):
h_start, h_end, w_start, w_end = view
local_states = norm_hidden_states_[:, h_start:h_end, w_start:w_end, :]
local_states = rearrange(local_states, 'bh h w d -> bh (h w) d')
@@ -175,7 +175,7 @@ def scale_forward(
norm_hidden_states_ = F.pad(norm_hidden_states, (0, 0, w_jitter_range, w_jitter_range, h_jitter_range, h_jitter_range), 'constant', 0)
value = torch.zeros_like(norm_hidden_states_)
count = torch.zeros_like(norm_hidden_states_)
- for index, view in enumerate(views):
+ for _index, view in enumerate(views):
h_start, h_end, w_start, w_end = view
local_states = norm_hidden_states_[:, h_start:h_end, w_start:w_end, :]
local_states = rearrange(local_states, 'bh h w d -> bh (h w) d')
@@ -198,7 +198,7 @@ def scale_forward(
value = torch.zeros_like(norm_hidden_states)
count = torch.zeros_like(norm_hidden_states)
- for index, global_view in enumerate(global_views):
+ for _index, global_view in enumerate(global_views):
h, w = global_view
global_states = norm_hidden_states[:, h::current_scale_num_h, w::current_scale_num_w, :]
global_states = rearrange(global_states, 'bh h w d -> bh (h w) d')
@@ -288,7 +288,7 @@ def ori_forward(
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
timestep: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
+ cross_attention_kwargs: Dict[str, Any] | None = None,
class_labels: Optional[torch.LongTensor] = None,
):
# Notice that normalization is always applied before the real computation in the following blocks.
diff --git a/scripts/infiniteyou/resampler.py b/scripts/infiniteyou/resampler.py
index 6d0011e83..f6f3d5c56 100644
--- a/scripts/infiniteyou/resampler.py
+++ b/scripts/infiniteyou/resampler.py
@@ -18,7 +18,7 @@ def FeedForward(dim, mult=4):
def reshape_tensor(x, heads):
- bs, length, width = x.shape
+ bs, length, _width = x.shape
#(bs, length, width) --> (bs, length, n_heads, dim_per_head)
x = x.view(bs, length, heads, -1)
# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
diff --git a/scripts/instantir/aggregator.py b/scripts/instantir/aggregator.py
index 5877340ca..63c216c39 100644
--- a/scripts/instantir/aggregator.py
+++ b/scripts/instantir/aggregator.py
@@ -113,9 +113,9 @@ class AggregatorOutput(BaseOutput):
class ConditioningEmbedding(nn.Module):
"""
Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN
- [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized
- training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the
- convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides
+ [11] to convert the entire dataset of 512 x 512 images into smaller 64 x 64 "latent images" for stabilized
+ training. This requires ControlNets to convert image-based conditions to 64 x 64 feature space to match the
+ convolution size. We use a tiny network E(·) of four convolution layers with 4 x 4 kernels and 2 x 2 strides
(activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full
model) to encode image-space conditions ... into feature maps ..."
"""
diff --git a/scripts/instantir/ip_adapter/attention_processor.py b/scripts/instantir/ip_adapter/attention_processor.py
index 191f45526..2b1a629a4 100644
--- a/scripts/instantir/ip_adapter/attention_processor.py
+++ b/scripts/instantir/ip_adapter/attention_processor.py
@@ -4,7 +4,7 @@ import torch.nn as nn
import torch.nn.functional as F
class AdaLayerNorm(nn.Module):
- def __init__(self, embedding_dim: int, time_embedding_dim: int = None):
+ def __init__(self, embedding_dim: int, time_embedding_dim: int | None = None):
super().__init__()
if time_embedding_dim is None:
@@ -219,7 +219,7 @@ class TA_IPAttnProcessor(nn.Module):
The context length of the image features.
"""
- def __init__(self, hidden_size, cross_attention_dim=None, time_embedding_dim: int = None, scale=1.0, num_tokens=4):
+ def __init__(self, hidden_size, cross_attention_dim=None, time_embedding_dim: int | None = None, scale=1.0, num_tokens=4):
super().__init__()
self.hidden_size = hidden_size
@@ -703,9 +703,9 @@ class AdditiveKV_AttnProcessor2_0(torch.nn.Module):
def __init__(
self,
- hidden_size: int = None,
- cross_attention_dim: int = None,
- time_embedding_dim: int = None,
+ hidden_size: int | None = None,
+ cross_attention_dim: int | None = None,
+ time_embedding_dim: int | None = None,
additive_scale: float = 1.0,
):
super().__init__()
@@ -812,9 +812,9 @@ class TA_AdditiveKV_AttnProcessor2_0(torch.nn.Module):
def __init__(
self,
- hidden_size: int = None,
- cross_attention_dim: int = None,
- time_embedding_dim: int = None,
+ hidden_size: int | None = None,
+ cross_attention_dim: int | None = None,
+ time_embedding_dim: int | None = None,
additive_scale: float = 1.0,
):
super().__init__()
@@ -969,7 +969,7 @@ class IPAttnProcessor2_0(torch.nn.Module):
if isinstance(encoder_hidden_states, tuple):
# FIXME: now hard coded to single image prompt.
- batch_size, _, hid_dim = encoder_hidden_states[0].shape
+ batch_size, _, _hid_dim = encoder_hidden_states[0].shape
ip_tokens = encoder_hidden_states[1][0]
encoder_hidden_states = torch.cat([encoder_hidden_states[0], ip_tokens], dim=1)
@@ -1066,7 +1066,7 @@ class TA_IPAttnProcessor2_0(torch.nn.Module):
The context length of the image features.
"""
- def __init__(self, hidden_size, cross_attention_dim=None, time_embedding_dim: int = None, scale=1.0, num_tokens=4):
+ def __init__(self, hidden_size, cross_attention_dim=None, time_embedding_dim: int | None = None, scale=1.0, num_tokens=4):
super().__init__()
if not hasattr(F, "scaled_dot_product_attention"):
@@ -1113,7 +1113,7 @@ class TA_IPAttnProcessor2_0(torch.nn.Module):
)
else:
# FIXME: now hard coded to single image prompt.
- batch_size, _, hid_dim = encoder_hidden_states[0].shape
+ batch_size, _, _hid_dim = encoder_hidden_states[0].shape
ip_hidden_states = encoder_hidden_states[1][0]
encoder_hidden_states = encoder_hidden_states[0]
batch_size, sequence_length, _ = (
@@ -1406,7 +1406,7 @@ def init_attn_proc(unet, ip_adapter_tokens=16, use_lcm=False, use_adaln=True, us
def init_aggregator_attn_proc(unet, use_adaln=False, split_attn=False):
attn_procs = {}
- unet_sd = unet.state_dict()
+ unet.state_dict()
for name in unet.attn_processors.keys():
# get layer name and hidden dim
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
diff --git a/scripts/instantir/ip_adapter/ip_adapter.py b/scripts/instantir/ip_adapter/ip_adapter.py
index 8a5df0a23..c8f853588 100644
--- a/scripts/instantir/ip_adapter/ip_adapter.py
+++ b/scripts/instantir/ip_adapter/ip_adapter.py
@@ -156,7 +156,7 @@ class IPAdapterPlus(torch.nn.Module):
orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj.parameters()]))
orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.ip_adapter.parameters()]))
org_unet_sum = []
- for attn_name, attn_proc in self.unet.attn_processors.items():
+ for _attn_name, attn_proc in self.unet.attn_processors.items():
if isinstance(attn_proc, (TA_IPAttnProcessor, IPAttnProcessor)):
org_unet_sum.append(torch.sum(torch.stack([torch.sum(p) for p in attn_proc.parameters()])))
org_unet_sum = torch.sum(torch.stack(org_unet_sum))
@@ -190,7 +190,7 @@ class IPAdapterPlus(torch.nn.Module):
# Verify if the weights loaded to unet
unet_sum = []
- for attn_name, attn_proc in self.unet.attn_processors.items():
+ for _attn_name, attn_proc in self.unet.attn_processors.items():
if isinstance(attn_proc, (TA_IPAttnProcessor, IPAttnProcessor)):
unet_sum.append(torch.sum(torch.stack([torch.sum(p) for p in attn_proc.parameters()])))
unet_sum = torch.sum(torch.stack(unet_sum))
diff --git a/scripts/instantir/ip_adapter/utils.py b/scripts/instantir/ip_adapter/utils.py
index 4f6d7075c..ed3411eeb 100644
--- a/scripts/instantir/ip_adapter/utils.py
+++ b/scripts/instantir/ip_adapter/utils.py
@@ -53,7 +53,7 @@ def init_adapter_in_unet(
if pretrained_model_path_or_dict is not None:
if "ip_adapter" in state_dict.keys():
adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
- missing, unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
+ missing, _unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
for mk in missing:
if "ln" not in mk:
raise ValueError(f"Missing keys in adapter_modules: {missing}")
@@ -143,7 +143,7 @@ def load_adapter_to_pipe(
# Load pretrinaed model if needed.
if "ip_adapter" in state_dict.keys():
adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
- missing, unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
+ missing, _unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
for mk in missing:
if "ln" not in mk:
raise ValueError(f"Missing keys in adapter_modules: {missing}")
diff --git a/scripts/instantir/lcm_single_step_scheduler.py b/scripts/instantir/lcm_single_step_scheduler.py
index a32affdc2..98dbc8732 100644
--- a/scripts/instantir/lcm_single_step_scheduler.py
+++ b/scripts/instantir/lcm_single_step_scheduler.py
@@ -322,7 +322,7 @@ class LCMSingleStepScheduler(SchedulerMixin, ConfigMixin):
def set_timesteps(
self,
- num_inference_steps: int = None,
+ num_inference_steps: int | None = None,
device: Union[str, torch.device] = None,
original_inference_steps: Optional[int] = None,
strength: int = 1.0,
diff --git a/scripts/instantir/sdxl_instantir.py b/scripts/instantir/sdxl_instantir.py
index bfb4f84e9..20beeb97c 100644
--- a/scripts/instantir/sdxl_instantir.py
+++ b/scripts/instantir/sdxl_instantir.py
@@ -166,11 +166,11 @@ PREVIEWER_LORA_MODULES = [
def remove_attn2(model):
def recursive_find_module(name, module):
- if not "up_blocks" in name and not "down_blocks" in name and not "mid_block" in name: return
- elif "resnets" in name: return
+ if "up_blocks" not in name and "down_blocks" not in name and "mid_block" not in name: return # noqa: E701
+ elif "resnets" in name: return # noqa: E701
if hasattr(module, "attn2"):
- setattr(module, "attn2", None)
- setattr(module, "norm2", None)
+ module.attn2 = None
+ module.norm2 = None
return
for sub_name, sub_module in module.named_children():
recursive_find_module(f"{name}.{sub_name}", sub_module)
@@ -388,7 +388,7 @@ class InstantIRPipeline(
if incompatible_keys is not None:
# check only for unexpected keys
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
- missing_keys = getattr(incompatible_keys, "missing_keys", None)
+ getattr(incompatible_keys, "missing_keys", None)
if unexpected_keys:
raise ValueError(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
@@ -834,12 +834,12 @@ class InstantIRPipeline(
)
if (
isinstance(self.aggregator, Aggregator)
- or is_compiled
- and isinstance(self.aggregator._orig_mod, Aggregator)
+ or (is_compiled
+ and isinstance(self.aggregator._orig_mod, Aggregator))
):
self.check_image(image, prompt, prompt_embeds)
else:
- assert False
+ raise AssertionError
if control_guidance_start >= control_guidance_end:
raise ValueError(
@@ -1064,13 +1064,13 @@ class InstantIRPipeline(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: Union[str, List[str]] | None = None,
prompt_2: Optional[Union[str, List[str]]] = None,
image: PipelineImageInput = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 30,
- timesteps: List[int] = None,
+ timesteps: List[int] | None = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 7.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
@@ -1098,9 +1098,9 @@ class InstantIRPipeline(
control_guidance_end: float = 1.0,
preview_start: float = 0.0,
preview_end: float = 1.0,
- original_size: Tuple[int, int] = None,
+ original_size: Tuple[int, int] | None = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ target_size: Tuple[int, int] | None = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
diff --git a/scripts/lbm/base/model_config.py b/scripts/lbm/base/model_config.py
index 427b69976..a956d907f 100644
--- a/scripts/lbm/base/model_config.py
+++ b/scripts/lbm/base/model_config.py
@@ -1,5 +1,5 @@
from pydantic.dataclasses import dataclass
-from ..config import BaseConfig
+from lbm.config import BaseConfig
@dataclass
diff --git a/scripts/lbm/config.py b/scripts/lbm/config.py
index de9948158..c9321db24 100644
--- a/scripts/lbm/config.py
+++ b/scripts/lbm/config.py
@@ -49,10 +49,10 @@ class BaseConfig:
f"Catch Exception {type(e)} with message: " + str(e)
) from e
- except FileNotFoundError:
+ except FileNotFoundError as e:
raise FileNotFoundError(
f"Config file not found. Please check path '{json_path}'"
- )
+ ) from e
@classmethod
def from_json(cls, json_path: str) -> "BaseConfig":
@@ -72,7 +72,7 @@ class BaseConfig:
warnings.warn(
f"You are trying to load a "
f"`{ cls.__name__}` while a "
- f"`{config_name}` is given."
+ f"`{config_name}` is given.", stacklevel=2
)
return cls.from_dict(config_dict)
@@ -134,7 +134,7 @@ class BaseConfig:
warnings.warn(
f"You are trying to load a "
f"`{ cls.__name__}` while a "
- f"`{config_name}` is given."
+ f"`{config_name}` is given.", stacklevel=2
)
return cls.from_dict(config_dict)
diff --git a/scripts/lbm/embedders/base/base_conditioner.py b/scripts/lbm/embedders/base/base_conditioner.py
index 4f6893902..c02a78af8 100644
--- a/scripts/lbm/embedders/base/base_conditioner.py
+++ b/scripts/lbm/embedders/base/base_conditioner.py
@@ -1,5 +1,5 @@
from typing import Any, Dict
-from ...base.base_model import BaseModel
+from lbm.base.base_model import BaseModel
from .base_conditioner_config import BaseConditionerConfig
diff --git a/scripts/lbm/embedders/base/base_conditioner_config.py b/scripts/lbm/embedders/base/base_conditioner_config.py
index 5a2eab0ee..fa23f4710 100644
--- a/scripts/lbm/embedders/base/base_conditioner_config.py
+++ b/scripts/lbm/embedders/base/base_conditioner_config.py
@@ -1,5 +1,5 @@
from pydantic.dataclasses import dataclass
-from ...config import BaseConfig
+from lbm.config import BaseConfig
@dataclass
diff --git a/scripts/lbm/embedders/conditioners_wrapper.py b/scripts/lbm/embedders/conditioners_wrapper.py
index de184153a..0b6fb2e75 100644
--- a/scripts/lbm/embedders/conditioners_wrapper.py
+++ b/scripts/lbm/embedders/conditioners_wrapper.py
@@ -35,7 +35,7 @@ class ConditionerWrapper(nn.Module):
for conditioner in self.conditioners:
cond_input_keys.append(conditioner.input_key)
- assert all([key in set(cond_input_keys) for key in self.ucg_keys])
+ assert all(key in set(cond_input_keys) for key in self.ucg_keys)
def on_fit_start(self, device: torch.device = None, *args, **kwargs):
for conditioner in self.conditioners:
@@ -44,7 +44,7 @@ class ConditionerWrapper(nn.Module):
def forward(
self,
batch: Dict[str, Any],
- ucg_keys: List[str] = None,
+ ucg_keys: List[str] | None = None,
set_ucg_rate_zero=False,
*args,
**kwargs,
diff --git a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py
index 0a60678c9..0dabaea8b 100644
--- a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py
+++ b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py
@@ -1,7 +1,7 @@
from dataclasses import field
from typing import List, Union
from pydantic.dataclasses import dataclass
-from ..base import BaseConditionerConfig
+from lbm.embedders.base import BaseConditionerConfig
@dataclass
diff --git a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py
index 447a6f8c0..2d1399466 100644
--- a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py
+++ b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py
@@ -1,8 +1,8 @@
from typing import Any, Dict
import torch
import torchvision.transforms.functional as F
-from ...vae import AutoencoderKLDiffusers
-from ..base import BaseConditioner
+from lbm.vae import AutoencoderKLDiffusers
+from lbm.embedders.base import BaseConditioner
from .latents_concat_embedder_config import LatentsConcatEmbedderConfig
diff --git a/scripts/lbm/lbm/lbm_config.py b/scripts/lbm/lbm/lbm_config.py
index a9b1b40ac..45fee17f6 100644
--- a/scripts/lbm/lbm/lbm_config.py
+++ b/scripts/lbm/lbm/lbm_config.py
@@ -1,6 +1,6 @@
from typing import List, Literal, Optional
from pydantic.dataclasses import dataclass
-from ..base import ModelConfig
+from lbm.base import ModelConfig
@dataclass
diff --git a/scripts/lbm/lbm/lbm_model.py b/scripts/lbm/lbm/lbm_model.py
index ac4f63332..e788adea3 100644
--- a/scripts/lbm/lbm/lbm_model.py
+++ b/scripts/lbm/lbm/lbm_model.py
@@ -5,10 +5,10 @@ import torch
import torch.nn as nn
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from tqdm import tqdm
-from ..base.base_model import BaseModel
-from ..embedders import ConditionerWrapper
-from ..unets import DiffusersUNet2DCondWrapper, DiffusersUNet2DWrapper
-from ..vae import AutoencoderKLDiffusers
+from lbm.base.base_model import BaseModel
+from lbm.embedders import ConditionerWrapper
+from lbm.unets import DiffusersUNet2DCondWrapper, DiffusersUNet2DWrapper
+from lbm.vae import AutoencoderKLDiffusers
from .lbm_config import LBMConfig
@@ -304,7 +304,7 @@ class LBMModel(BaseModel):
def _get_conditioning(
self,
batch: Dict[str, Any],
- ucg_keys: List[str] = None,
+ ucg_keys: List[str] | None = None,
set_ucg_rate_zero=False,
*args,
**kwargs,
diff --git a/scripts/lbm/tiler.py b/scripts/lbm/tiler.py
index 529348332..d32f82dd8 100644
--- a/scripts/lbm/tiler.py
+++ b/scripts/lbm/tiler.py
@@ -17,7 +17,7 @@ class Tiler:
overlap_size: tuple,
scale: int = 1,
out_channels: int = 3,
- ) -> List[List[torch.tensor]]:
+ ) -> List[List[torch.Tensor]]:
"""Get tiles
Args:
input (torch.Tensor): input array of shape (batch_size, channels, height, width)
@@ -82,8 +82,8 @@ class Tiler:
return tiles
def merge_tiles(
- self, tiles: List[List[torch.tensor]], tiling_method: str = "gaussian"
- ) -> torch.tensor:
+ self, tiles: List[List[torch.Tensor]], tiling_method: str = "gaussian"
+ ) -> torch.Tensor:
"""Merge tiles by averaging the overlaping regions
Args:
tiles (Dict[str, Tile]): dictionary of processed tiles
@@ -102,7 +102,7 @@ class Tiler:
f"Unknown tiling method {tiling_method}. Available methods are {TILING_METHODS}"
)
- def _average_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.tensor:
+ def _average_merge_tiles(self, tiles: List[List[torch.Tensor]]) -> torch.Tensor:
"""Merge tiles by averaging the overlaping regions
Args:
tiles (Dict[str, Tile]): dictionary of processed tiles
@@ -203,10 +203,10 @@ class Tiler:
torch.tensor(weights, device="cpu"), (nbatches, channels, 1, 1)
)
- def _gaussian_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.tensor:
+ def _gaussian_merge_tiles(self, tiles: List[List[torch.Tensor]]) -> torch.Tensor:
"""Merge tiles by averaging the overlaping regions
Args:
- List[List[torch.tensor]]: List of processed tiles
+ List[List[torch.Tensor]]: List of processed tiles
Returns:
torch.tensor: output image
"""
@@ -277,10 +277,10 @@ class Tiler:
] * (x / blend_extent)
return b
- def _linear_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.Tensor:
+ def _linear_merge_tiles(self, tiles: List[List[torch.Tensor]]) -> torch.Tensor:
"""Merge tiles by blending the overlaping regions
Args:
- tiles (List[List[torch.tensor]]): List of processed tiles
+ tiles (List[List[torch.Tensor]]): List of processed tiles
Returns:
torch.Tensor: output image
"""
diff --git a/scripts/lbm/unets/unet.py b/scripts/lbm/unets/unet.py
index 6aa6d7eee..8fddb5332 100644
--- a/scripts/lbm/unets/unet.py
+++ b/scripts/lbm/unets/unet.py
@@ -17,7 +17,7 @@ class DiffusersUNet2DWrapper(UNet2DModel):
self,
sample: torch.Tensor,
timestep: Union[torch.Tensor, float, int],
- conditioning: Dict[str, torch.Tensor] = None,
+ conditioning: Dict[str, torch.Tensor] | None = None,
*args,
**kwargs,
):
diff --git a/scripts/lbm/vae/autoencoderKL.py b/scripts/lbm/vae/autoencoderKL.py
index 09a5523bb..609156bc9 100644
--- a/scripts/lbm/vae/autoencoderKL.py
+++ b/scripts/lbm/vae/autoencoderKL.py
@@ -1,7 +1,7 @@
import torch
from diffusers.models import AutoencoderKL
-from ..base.base_model import BaseModel
-from ..tiler import Tiler, pad
+from lbm.base.base_model import BaseModel
+from lbm.tiler import Tiler, pad
from .autoencoderKL_config import AutoencoderKLDiffusersConfig
@@ -56,7 +56,7 @@ class AutoencoderKLDiffusers(BaseModel):
# set downsampling factor
self.downsampling_factor = int(x.shape[2] / z.shape[2])
- def encode(self, x: torch.tensor, batch_size: int = 8):
+ def encode(self, x: torch.Tensor, batch_size: int = 8):
latents = []
for i in range(0, x.shape[0], batch_size):
latents.append(
@@ -67,7 +67,7 @@ class AutoencoderKLDiffusers(BaseModel):
return latents
- def decode(self, z: torch.tensor):
+ def decode(self, z: torch.Tensor):
if self.has_latents_mean and self.has_latents_std:
latents_mean = (
diff --git a/scripts/lbm/vae/autoencoderKL_config.py b/scripts/lbm/vae/autoencoderKL_config.py
index a2f1ed802..2fe4aee65 100644
--- a/scripts/lbm/vae/autoencoderKL_config.py
+++ b/scripts/lbm/vae/autoencoderKL_config.py
@@ -1,6 +1,6 @@
from typing import Tuple
from pydantic.dataclasses import dataclass
-from ..base import ModelConfig
+from lbm.base import ModelConfig
@dataclass
diff --git a/scripts/nudenet/imageguard.py b/scripts/nudenet/imageguard.py
index c707402b2..5976527fa 100644
--- a/scripts/nudenet/imageguard.py
+++ b/scripts/nudenet/imageguard.py
@@ -150,8 +150,8 @@ def image_guard(image, policy:str | None=None, model_name:str=''):
prompt = processor.apply_chat_template(chat_template, add_generation_prompt=True)
inputs = processor(text=prompt, images=image, return_tensors="pt")
input_ids = inputs.input_ids
- model = model.to(device=devices.device)
- inputs = {k: v.to(device=devices.device) for k, v in inputs.items()}
+ model = model.to(devices.device)
+ inputs = {k: v.to(devices.device) for k, v in inputs.items()}
kwargs = {
"max_new_tokens": 200,
"do_sample": True,
@@ -162,7 +162,7 @@ def image_guard(image, policy:str | None=None, model_name:str=''):
"use_cache": True,
}
generated_ids = model.generate(**inputs, **kwargs)
- model = model.to(device=devices.cpu)
+ model = model.to(devices.cpu)
trimmed_ids = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(input_ids, generated_ids)]
# output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
diff --git a/scripts/nudenet/langdetect.py b/scripts/nudenet/langdetect.py
index 0606d841e..ea8652a18 100644
--- a/scripts/nudenet/langdetect.py
+++ b/scripts/nudenet/langdetect.py
@@ -3,7 +3,7 @@ repo_id = "facebook/fasttext-language-identification"
model = None
-def lang_detect(text:str, top:int=1, threshold:float=0.25) -> str:
+def lang_detect(text:str, top:int=1, threshold:float=0.25) -> str | list[str]:
try:
global model # pylint: disable=global-statement
from modules import shared
diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py
index 764b978ac..3162f0787 100644
--- a/scripts/prompt_enhance.py
+++ b/scripts/prompt_enhance.py
@@ -369,7 +369,7 @@ class PromptEnhanceScript(scripts_manager.Script):
sd_models.move_model(self.llm, devices.cpu, force=True)
self.llm = None
self.tokenizer = None
- devices.torch_gc(force=True, reason='prompt enhance model switch')
+ devices.torch_gc(force=True, reason='prompt-enhance:load')
log.debug(f'Prompt enhance: name="{self.model}" unload')
self.model = None
load_args = { 'pretrained_model_name_or_path': model_repo if not gguf_args else model_gguf }
@@ -439,7 +439,7 @@ class PromptEnhanceScript(scripts_manager.Script):
self.model = None
self.llm = None
self.tokenizer = None
- devices.torch_gc(force=True, reason='prompt enhance unload')
+ devices.torch_gc(force=True, reason='prompt-enhance:unload')
log.debug(f'Prompt enhance: model="{model_name}" unloaded')
else:
log.debug('Prompt enhance: no model loaded')
@@ -801,7 +801,7 @@ class PromptEnhanceScript(scripts_manager.Script):
response = f'Error: {str(e)}'
finally:
offload_aux('prompt_enhance')
- devices.torch_gc(force=True, reason='prompt enhance offload')
+ devices.torch_gc(force=False, reason='prompt-enhance')
t1 = time.time()
if isinstance(response, list):
@@ -958,6 +958,8 @@ class PromptEnhanceScript(scripts_manager.Script):
shared.prompt_styles.apply_styles_to_extra(p)
p.styles = []
jobid = shared.state.begin('LLM')
+ p.extra_generation_params['LLM'] = get_model_repo_from_display(llm_model)
+ p.extra_generation_params['Original'] = p.prompt
p.prompt = self.enhance(
prompt=p.prompt,
seed=p.seed,
@@ -980,5 +982,4 @@ class PromptEnhanceScript(scripts_manager.Script):
keep_thinking=keep_thinking,
)
timer.process.record('prompt')
- p.extra_generation_params['LLM'] = llm_model
shared.state.end(jobid)
diff --git a/scripts/pulid/eva_clip/eva_vit_model.py b/scripts/pulid/eva_clip/eva_vit_model.py
index 032500241..61fe1cfe7 100644
--- a/scripts/pulid/eva_clip/eva_vit_model.py
+++ b/scripts/pulid/eva_clip/eva_vit_model.py
@@ -171,7 +171,7 @@ class Attention(nn.Module):
self.rope = rope
def forward(self, x, rel_pos_bias=None, attn_mask=None):
- B, N, C = x.shape
+ B, N, _C = x.shape
if self.subln:
q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
@@ -318,7 +318,7 @@ class PatchEmbed(nn.Module):
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x, **kwargs):
- B, C, H, W = x.shape
+ _B, _C, H, W = x.shape
# FIXME look at relaxing size constraints
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
@@ -493,7 +493,7 @@ class EVAVisionTransformer(nn.Module):
def forward_features(self, x, return_all_features=False, return_hidden=False, shuffle=False):
x = self.patch_embed(x)
- batch_size, seq_len, _ = x.size()
+ batch_size, _seq_len, _ = x.size()
if shuffle:
idx = torch.randperm(x.shape[1]) + 1
diff --git a/scripts/pulid/eva_clip/factory.py b/scripts/pulid/eva_clip/factory.py
index b33929625..3a466d605 100644
--- a/scripts/pulid/eva_clip/factory.py
+++ b/scripts/pulid/eva_clip/factory.py
@@ -18,7 +18,7 @@ from .tokenizer import HFTokenizer, tokenize
from .utils import resize_clip_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed, resize_eva_pos_embed
-_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / "model_configs/"]
_MODEL_CONFIGS = {} # directory (model_name: config) of model architecture configs
@@ -220,8 +220,8 @@ def create_model(
pretrained_image: str = '',
pretrained_text: str = '',
pretrained_hf: bool = True,
- pretrained_visual_model: str = None,
- pretrained_text_model: str = None,
+ pretrained_visual_model: str | None = None,
+ pretrained_text_model: str | None = None,
cache_dir: Optional[str] = None,
skip_list: list = [],
):
@@ -367,8 +367,8 @@ def create_model_and_transforms(
pretrained_image: str = '',
pretrained_text: str = '',
pretrained_hf: bool = True,
- pretrained_visual_model: str = None,
- pretrained_text_model: str = None,
+ pretrained_visual_model: str | None = None,
+ pretrained_text_model: str | None = None,
image_mean: Optional[Tuple[float, ...]] = None,
image_std: Optional[Tuple[float, ...]] = None,
cache_dir: Optional[str] = None,
@@ -422,8 +422,8 @@ def create_transforms(
pretrained_image: str = '',
pretrained_text: str = '',
pretrained_hf: bool = True,
- pretrained_visual_model: str = None,
- pretrained_text_model: str = None,
+ pretrained_visual_model: str | None = None,
+ pretrained_text_model: str | None = None,
image_mean: Optional[Tuple[float, ...]] = None,
image_std: Optional[Tuple[float, ...]] = None,
cache_dir: Optional[str] = None,
diff --git a/scripts/pulid/eva_clip/hf_model.py b/scripts/pulid/eva_clip/hf_model.py
index 1665ada0b..275b9fe32 100644
--- a/scripts/pulid/eva_clip/hf_model.py
+++ b/scripts/pulid/eva_clip/hf_model.py
@@ -14,7 +14,7 @@ try:
from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer, AutoConfig, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
BaseModelOutputWithPoolingAndCrossAttentions
-except ImportError as e:
+except ImportError:
transformers = None
@@ -77,10 +77,10 @@ class HFTextEncoder(nn.Module):
self,
model_name_or_path: str,
output_dim: int,
- tokenizer_name: str = None,
+ tokenizer_name: str | None = None,
config: PretrainedConfig = None,
- pooler_type: str = None,
- proj: str = None,
+ pooler_type: str | None = None,
+ proj: str | None = None,
pretrained: bool = True,
masked_language_modeling: bool = False):
super().__init__()
diff --git a/scripts/pulid/eva_clip/model.py b/scripts/pulid/eva_clip/model.py
index 05b055794..5be65b752 100644
--- a/scripts/pulid/eva_clip/model.py
+++ b/scripts/pulid/eva_clip/model.py
@@ -385,7 +385,7 @@ def build_model_from_openai_state_dict(
vocab_size = state_dict["token_embedding.weight"].shape[0]
transformer_width = state_dict["ln_final.weight"].shape[0]
transformer_heads = transformer_width // 64
- transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+ transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
vision_cfg = CLIPVisionCfg(
layers=vision_layers,
diff --git a/scripts/pulid/eva_clip/modified_resnet.py b/scripts/pulid/eva_clip/modified_resnet.py
index 151bfdd0b..09e7ad019 100644
--- a/scripts/pulid/eva_clip/modified_resnet.py
+++ b/scripts/pulid/eva_clip/modified_resnet.py
@@ -56,7 +56,7 @@ class Bottleneck(nn.Module):
class AttentionPool2d(nn.Module):
- def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+ def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int | None = None):
super().__init__()
self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
self.k_proj = nn.Linear(embed_dim, embed_dim)
diff --git a/scripts/pulid/eva_clip/openai.py b/scripts/pulid/eva_clip/openai.py
index cc4e13e87..73f21fe0b 100644
--- a/scripts/pulid/eva_clip/openai.py
+++ b/scripts/pulid/eva_clip/openai.py
@@ -68,7 +68,7 @@ def load_openai_model(
except RuntimeError:
# loading saved state dict
if jit:
- warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+ warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead", stacklevel=2)
jit = False
state_dict = torch.load(model_path, map_location="cpu")
diff --git a/scripts/pulid/eva_clip/pretrained.py b/scripts/pulid/eva_clip/pretrained.py
index bb87c540c..bf957db74 100644
--- a/scripts/pulid/eva_clip/pretrained.py
+++ b/scripts/pulid/eva_clip/pretrained.py
@@ -261,7 +261,7 @@ def download_pretrained_from_url(
if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
return download_target
else:
- warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+ warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file", stacklevel=2)
else:
return download_target
diff --git a/scripts/pulid/eva_clip/rope.py b/scripts/pulid/eva_clip/rope.py
index 764cb1976..6fc6cec2f 100644
--- a/scripts/pulid/eva_clip/rope.py
+++ b/scripts/pulid/eva_clip/rope.py
@@ -6,7 +6,7 @@ import logging
def broadcat(tensors, dim = -1):
num_tensors = len(tensors)
- shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+ shape_lens = set(map(lambda t: len(t.shape), tensors))
assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
shape_len = list(shape_lens)[0]
dim = (dim + shape_len) if dim < 0 else dim
@@ -51,7 +51,7 @@ class VisionRotaryEmbedding(nn.Module):
else:
raise ValueError(f'unknown modality {freqs_for}')
- if ft_seq_len is None: ft_seq_len = pt_seq_len
+ if ft_seq_len is None: ft_seq_len = pt_seq_len # noqa: E701
t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
freqs_h = torch.einsum('..., f -> ... f', t, freqs)
@@ -101,7 +101,7 @@ class VisionRotaryEmbeddingFast(nn.Module):
else:
raise ValueError(f'unknown modality {freqs_for}')
- if ft_seq_len is None: ft_seq_len = pt_seq_len
+ if ft_seq_len is None: ft_seq_len = pt_seq_len # noqa: E701
t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
freqs = torch.einsum('..., f -> ... f', t, freqs)
diff --git a/scripts/pulid/eva_clip/timm_model.py b/scripts/pulid/eva_clip/timm_model.py
index 53bc4d469..508cf24c3 100644
--- a/scripts/pulid/eva_clip/timm_model.py
+++ b/scripts/pulid/eva_clip/timm_model.py
@@ -91,8 +91,8 @@ class TimmModel(nn.Module):
try:
# FIXME import here until API stable and in an official release
from timm.models.helpers import group_parameters, group_modules
- except ImportError:
- raise RuntimeError('Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
+ except ImportError as e:
+ raise RuntimeError('Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`') from e
matcher = self.trunk.group_matcher()
gparams = group_parameters(self.trunk, matcher)
max_layer_id = max(gparams.keys())
@@ -110,7 +110,7 @@ class TimmModel(nn.Module):
def set_grad_checkpointing(self, enable=True):
try:
self.trunk.set_grad_checkpointing(enable)
- except Exception as e:
+ except Exception:
logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
def forward(self, x):
diff --git a/scripts/pulid/eva_clip/tokenizer.py b/scripts/pulid/eva_clip/tokenizer.py
index b76e2a3aa..45fa860a9 100644
--- a/scripts/pulid/eva_clip/tokenizer.py
+++ b/scripts/pulid/eva_clip/tokenizer.py
@@ -12,7 +12,6 @@ import regex as re
import torch
# https://stackoverflow.com/q/62691279
-import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
diff --git a/scripts/pulid/eva_clip/transform.py b/scripts/pulid/eva_clip/transform.py
index f74b750d9..7419a00f2 100644
--- a/scripts/pulid/eva_clip/transform.py
+++ b/scripts/pulid/eva_clip/transform.py
@@ -18,7 +18,7 @@ class ResizeMaxSize(nn.Module):
raise TypeError(f"Size should be int. Got {type(max_size)}")
self.max_size = max_size
self.interpolation = interpolation
- self.fn = min if fn == 'min' else min
+ self.fn = min
self.fill = fill
def forward(self, img):
diff --git a/scripts/pulid/eva_clip/transformer.py b/scripts/pulid/eva_clip/transformer.py
index 57a24ad2c..e5322ef5c 100644
--- a/scripts/pulid/eva_clip/transformer.py
+++ b/scripts/pulid/eva_clip/transformer.py
@@ -279,8 +279,8 @@ class CustomAttention(nn.Module):
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
N_q, B_q, C_q = q.shape
- N_k, B_k, C_k = k.shape
- N_v, B_v, C_v = v.shape
+ N_k, B_k, _C_k = k.shape
+ N_v, B_v, _C_v = v.shape
if self.xattn:
# B, N, C -> B, N, num_heads, C
q = q.permute(1, 0, 2).reshape(B_q, N_q, self.num_heads, -1)
@@ -335,7 +335,7 @@ class CustomResidualAttentionBlock(nn.Module):
d_model: int,
n_head: int,
mlp_ratio: float = 4.0,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
act_layer: Callable = nn.GELU,
norm_layer: Callable = LayerNorm,
scale_cosine_attn: bool = False,
@@ -386,7 +386,7 @@ class CustomTransformer(nn.Module):
layers: int,
heads: int,
mlp_ratio: float = 4.0,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
act_layer: Callable = nn.GELU,
norm_layer: Callable = LayerNorm,
scale_cosine_attn: bool = True,
@@ -436,7 +436,7 @@ class ResidualAttentionBlock(nn.Module):
d_model: int,
n_head: int,
mlp_ratio: float = 4.0,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
act_layer: Callable = nn.GELU,
norm_layer: Callable = LayerNorm,
xattn: bool = False,
@@ -479,7 +479,7 @@ class Transformer(nn.Module):
layers: int,
heads: int,
mlp_ratio: float = 4.0,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
act_layer: Callable = nn.GELU,
norm_layer: Callable = LayerNorm,
xattn: bool = False,
@@ -513,7 +513,7 @@ class VisionTransformer(nn.Module):
layers: int,
heads: int,
mlp_ratio: float,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
patch_dropout: float = 0.,
global_average_pool: bool = False,
output_dim: int = 512,
@@ -634,7 +634,7 @@ class TextTransformer(nn.Module):
width: int = 512,
heads: int = 8,
layers: int = 12,
- ls_init_value: float = None,
+ ls_init_value: float | None = None,
output_dim: int = 512,
act_layer: Callable = nn.GELU,
norm_layer: Callable = LayerNorm,
diff --git a/scripts/pulid/eva_clip/utils.py b/scripts/pulid/eva_clip/utils.py
index 1c3c06201..398982438 100644
--- a/scripts/pulid/eva_clip/utils.py
+++ b/scripts/pulid/eva_clip/utils.py
@@ -76,7 +76,7 @@ def resize_visual_pos_embed(state_dict, model, interpolation: str = 'bicubic', s
state_dict['positional_embedding'] = new_pos_embed
def resize_evaclip_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
- all_keys = list(state_dict.keys())
+ list(state_dict.keys())
# interpolate position embedding
if 'visual.pos_embed' in state_dict:
pos_embed_checkpoint = state_dict['visual.pos_embed']
@@ -107,7 +107,7 @@ def resize_evaclip_pos_embed(state_dict, model, interpolation: str = 'bicubic',
def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
- all_keys = list(state_dict.keys())
+ list(state_dict.keys())
# interpolate position embedding
if 'pos_embed' in state_dict:
pos_embed_checkpoint = state_dict['pos_embed']
@@ -149,7 +149,7 @@ def resize_rel_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_
dst_num_pos, _ = model.visual.state_dict()[key].size()
dst_patch_shape = model.visual.patch_embed.patch_shape
if dst_patch_shape[0] != dst_patch_shape[1]:
- raise NotImplementedError()
+ raise NotImplementedError
num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
diff --git a/scripts/pulid/pulid_flux.py b/scripts/pulid/pulid_flux.py
index 13ea880c4..8e90d6b27 100644
--- a/scripts/pulid/pulid_flux.py
+++ b/scripts/pulid/pulid_flux.py
@@ -6,7 +6,7 @@ from modules.logger import log
def apply_flux(pipe: FluxPipeline):
- if not hasattr(pipe, 'transformer') or not 'Nunchaku' in pipe.transformer.__class__.__name__:
+ if not hasattr(pipe, 'transformer') or 'Nunchaku' not in pipe.transformer.__class__.__name__:
log.error('PuLID: flux support requires nunchaku')
return pipe
diff --git a/scripts/pulid/pulid_sdxl.py b/scripts/pulid/pulid_sdxl.py
index 2fc38373b..989edabe8 100644
--- a/scripts/pulid/pulid_sdxl.py
+++ b/scripts/pulid/pulid_sdxl.py
@@ -34,7 +34,7 @@ class StableDiffusionXLPuLIDPipeline:
pipe: Union[StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline],
device: torch.device,
dtype: torch.dtype=None,
- providers: list=None,
+ providers: list | None=None,
offload: bool=True,
sampler=None,
cache_dir=None,
@@ -284,7 +284,7 @@ class StableDiffusionXLPuLIDPipeline:
debug(f'PulID embedding: cond={id_embedding.shape} uncond={uncond_id_embedding.shape}')
return uncond_id_embedding, id_embedding
- def set_progress_bar_config(self, bar_format: str = None, ncols: int = 80, colour: str = None):
+ def set_progress_bar_config(self, bar_format: str | None = None, ncols: int = 80, colour: str | None = None):
import functools
from tqdm.auto import trange as trange_orig
import pulid_sampling
diff --git a/scripts/style_aligned_ext.py b/scripts/style_aligned_ext.py
index c64d811a8..6eac8d6ed 100644
--- a/scripts/style_aligned_ext.py
+++ b/scripts/style_aligned_ext.py
@@ -35,7 +35,7 @@ class StyleAlignedScript(scripts_manager.Script):
def ui(self, _is_img2img): # ui elements
with gr.Row():
- gr.HTML('  Style Aligned Image Generation ')
+ gr.HTML('  Style Aligned Image Generation ')
with gr.Row():
preset = gr.Dropdown(label="Preset", choices=['text', 'image', 'all'], value='text')
scheduler = gr.Checkbox(label="Override scheduler", value=False)
diff --git a/scripts/ultimate-upscale.py b/scripts/ultimate-upscale.py
new file mode 100644
index 000000000..f620945f7
--- /dev/null
+++ b/scripts/ultimate-upscale.py
@@ -0,0 +1,515 @@
+from enum import Enum
+import math
+import gradio as gr
+from PIL import Image, ImageDraw, ImageOps
+from modules import processing, shared, images, scripts_manager
+from modules.processing import StableDiffusionProcessing
+from modules.processing import Processed
+from modules.shared import opts, state
+from modules.logger import log
+
+elem_id = "ultimateupscale"
+
+class USDUMode(Enum):
+ LINEAR = 0
+ CHESS = 1
+ NONE = 2
+
+class USDUSFMode(Enum):
+ NONE = 0
+ BAND_PASS = 1
+ HALF_TILE = 2
+ HALF_TILE_PLUS_INTERSECTIONS = 3
+
+class USDUpscaler():
+
+ def __init__(self, p, image, upscaler_index:int, tile_width, tile_height) -> None:
+ self.p:StableDiffusionProcessing = p
+ self.image:Image.Image = image
+ self.scales = []
+ self.result_images = []
+ self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height))
+ self.upscaler = shared.sd_upscalers[upscaler_index]
+ self.redraw = USDURedraw()
+ self.redraw.tile_width = tile_width if tile_width > 0 else tile_height
+ self.redraw.tile_height = tile_height if tile_height > 0 else tile_width
+ self.seams_fix = USDUSeamsFix()
+ self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height
+ self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width
+ self.initial_info = None
+ self.rows = math.ceil(self.p.height / self.redraw.tile_height)
+ self.cols = math.ceil(self.p.width / self.redraw.tile_width)
+
+ def get_factor(self, num):
+ # Its just return, don't need elif
+ if num == 1:
+ return 2
+ if num % 4 == 0:
+ return 4
+ if num % 3 == 0:
+ return 3
+ if num % 2 == 0:
+ return 2
+ return 0
+
+ def get_factors(self):
+ scales = []
+ current_scale = 1
+ current_scale_factor = self.get_factor(self.scale_factor)
+ while current_scale_factor == 0:
+ self.scale_factor += 1
+ current_scale_factor = self.get_factor(self.scale_factor)
+ while current_scale < self.scale_factor:
+ current_scale_factor = self.get_factor(self.scale_factor // current_scale)
+ scales.append(current_scale_factor)
+ current_scale = current_scale * current_scale_factor
+ if current_scale_factor == 0:
+ break
+ self.scales = enumerate(scales)
+
+ def upscale(self):
+ # Log info
+ log.info(f'USDUpscaler: canvas={self.p.width}x{self.p.height} image={self.image.width}x{self.image.height} scale={self.scale_factor} scales={list(self.scales)} upscaler={self.upscaler.name}')
+ # Check upscaler is not empty
+ if self.upscaler.name == "None":
+ self.image = self.image.resize((self.p.width, self.p.height), resample=Image.Resampling.LANCZOS)
+ return
+ # Get list with scale factors
+ self.get_factors()
+ # Upscaling image over all factors
+ for _index, value in self.scales:
+ self.image = self.upscaler.scaler.upscale(self.image, value, self.upscaler.data_path)
+ # Resize image to set values
+ self.image = self.image.resize((self.p.width, self.p.height), resample=Image.Resampling.LANCZOS)
+
+ def setup_redraw(self, redraw_mode, padding, mask_blur):
+ self.redraw.mode = USDUMode(redraw_mode)
+ self.redraw.enabled = self.redraw.mode != USDUMode.NONE
+ self.redraw.padding = padding
+ self.p.mask_blur = mask_blur
+
+ def setup_seams_fix(self, padding, denoise, mask_blur, width, mode):
+ self.seams_fix.padding = padding
+ self.seams_fix.denoise = denoise
+ self.seams_fix.mask_blur = mask_blur
+ self.seams_fix.width = width
+ self.seams_fix.mode = USDUSFMode(mode)
+ self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE
+
+ def save_image(self):
+ if type(self.p.prompt) != list:
+ images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt, opts.samples_format, info=self.initial_info, p=self.p)
+ else:
+ images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt[0], opts.samples_format, info=self.initial_info, p=self.p)
+
+ def calc_jobs_count(self):
+ redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0
+ seams_job_count = 0
+ if self.seams_fix.mode == USDUSFMode.BAND_PASS:
+ seams_job_count = self.rows + self.cols - 2
+ elif self.seams_fix.mode == USDUSFMode.HALF_TILE:
+ seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
+ elif self.seams_fix.mode == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
+ seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + (self.rows - 1) * (self.cols - 1)
+ state.job_count = redraw_job_count + seams_job_count
+
+ def add_extra_info(self):
+ log.info(f"USDUpscaler: tile={self.redraw.tile_width}x{self.redraw.tile_height} rows={self.rows} cols={self.cols} redraw={self.redraw.enabled} seams={self.seams_fix.mode.name}")
+ self.p.extra_generation_params["USDUpscaler upscaler"] = self.upscaler.name
+ self.p.extra_generation_params["USDUpscaler tile"] = f'{self.redraw.tile_width}x{self.redraw.tile_height}'
+ self.p.extra_generation_params["USDUpscaler blur"] = self.p.mask_blur
+ self.p.extra_generation_params["USDUpscaler padding"] = self.redraw.padding
+
+ def process(self):
+ state.begin()
+ self.calc_jobs_count()
+ self.result_images = []
+ if self.redraw.enabled:
+ self.image = self.redraw.start(self.p, self.image, self.rows, self.cols)
+ self.initial_info = self.redraw.initial_info
+ self.result_images.append(self.image)
+ if self.seams_fix.enabled:
+ self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols)
+ self.initial_info = self.seams_fix.initial_info
+ self.result_images.append(self.image)
+ state.end()
+
+class USDURedraw():
+ def __init__(self) -> None:
+ self.enabled = False
+ self.mode = USDUMode.NONE
+ self.padding = 0
+ self.tile_width = 0
+ self.tile_height = 0
+ self.initial_info = None
+
+ def init_draw(self, p, width, height):
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = self.padding
+ p.width = math.ceil((self.tile_width+self.padding) / 64) * 64
+ p.height = math.ceil((self.tile_height+self.padding) / 64) * 64
+ mask = Image.new("L", (width, height), "black")
+ draw = ImageDraw.Draw(mask)
+ return mask, draw
+
+ def calc_rectangle(self, xi, yi):
+ x1 = xi * self.tile_width
+ y1 = yi * self.tile_height
+ x2 = xi * self.tile_width + self.tile_width
+ y2 = yi * self.tile_height + self.tile_height
+ return x1, y1, x2, y2
+
+ def linear_process(self, p, image, rows, cols):
+ mask, draw = self.init_draw(p, image.width, image.height)
+ processed = None
+ for yi in range(rows):
+ for xi in range(cols):
+ if state.interrupted:
+ break
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
+ if len(processed.images) > 0:
+ image = processed.images[0]
+ p.width = image.width
+ p.height = image.height
+ if processed is not None:
+ self.initial_info = processed.infotext(p, 0)
+ return image
+
+ def chess_process(self, p, image, rows, cols):
+ mask, draw = self.init_draw(p, image.width, image.height)
+ tiles = []
+ processed = None
+ # calc tiles colors
+ for yi in range(rows):
+ for xi in range(cols):
+ if state.interrupted:
+ break
+ if xi == 0:
+ tiles.append([])
+ color = xi % 2 == 0
+ if yi > 0 and yi % 2 != 0:
+ color = not color
+ tiles[yi].append(color)
+
+ for yi in range(len(tiles)):
+ for xi in range(len(tiles[yi])):
+ if state.interrupted:
+ break
+ if not tiles[yi][xi]:
+ tiles[yi][xi] = not tiles[yi][xi]
+ continue
+ tiles[yi][xi] = not tiles[yi][xi]
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
+ if len(processed.images) > 0:
+ image = processed.images[0]
+
+ for yi in range(len(tiles)):
+ for xi in range(len(tiles[yi])):
+ if state.interrupted:
+ break
+ if not tiles[yi][xi]:
+ continue
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
+ if len(processed.images) > 0:
+ image = processed.images[0]
+
+ p.width = image.width
+ p.height = image.height
+ if processed is not None:
+ self.initial_info = processed.infotext(p, 0)
+
+ return image
+
+ def start(self, p, image, rows, cols):
+ self.initial_info = None
+ if self.mode == USDUMode.LINEAR:
+ return self.linear_process(p, image, rows, cols)
+ if self.mode == USDUMode.CHESS:
+ return self.chess_process(p, image, rows, cols)
+ return image
+
+class USDUSeamsFix():
+ def __init__(self) -> None:
+ self.enabled = False
+ self.padding = 0
+ self.denoise = 0
+ self.mask_blur = 0
+ self.width = 0
+ self.mode = USDUSFMode.NONE
+ self.tile_width = 0
+ self.tile_height = 0
+ self.initial_info = None
+
+ def init_draw(self, p):
+ self.initial_info = None
+ p.width = math.ceil((self.tile_width+self.padding) / 64) * 64
+ p.height = math.ceil((self.tile_height+self.padding) / 64) * 64
+
+ def half_tile_process(self, p, image, rows, cols):
+ self.init_draw(p)
+ processed = None
+ gradient = Image.linear_gradient("L")
+ row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
+ row_gradient.paste(gradient.resize(
+ (self.tile_width, self.tile_height//2), resample=Image.Resampling.BICUBIC), (0, 0))
+ row_gradient.paste(gradient.rotate(180).resize(
+ (self.tile_width, self.tile_height//2), resample=Image.Resampling.BICUBIC),
+ (0, self.tile_height//2))
+ col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
+ col_gradient.paste(gradient.rotate(90).resize(
+ (self.tile_width//2, self.tile_height), resample=Image.Resampling.BICUBIC), (0, 0))
+ col_gradient.paste(gradient.rotate(270).resize(
+ (self.tile_width//2, self.tile_height), resample=Image.Resampling.BICUBIC), (self.tile_width//2, 0))
+
+ p.denoising_strength = self.denoise
+ p.mask_blur = self.mask_blur
+
+ for yi in range(rows-1):
+ for xi in range(cols):
+ if state.interrupted:
+ break
+ p.width = self.tile_width
+ p.height = self.tile_height
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = self.padding
+ mask = Image.new("L", (image.width, image.height), "black")
+ mask.paste(row_gradient, (xi*self.tile_width, yi*self.tile_height + self.tile_height//2))
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ if len(processed.images) > 0:
+ image = processed.images[0]
+
+ for yi in range(rows):
+ for xi in range(cols-1):
+ if state.interrupted:
+ break
+ p.width = self.tile_width
+ p.height = self.tile_height
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = self.padding
+ mask = Image.new("L", (image.width, image.height), "black")
+ mask.paste(col_gradient, (xi*self.tile_width+self.tile_width//2, yi*self.tile_height))
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ if len(processed.images) > 0:
+ image = processed.images[0]
+
+ p.width = image.width
+ p.height = image.height
+ if processed is not None:
+ self.initial_info = processed.infotext(p, 0)
+ return image
+
+ def half_tile_process_corners(self, p, image, rows, cols):
+ fixed_image = self.half_tile_process(p, image, rows, cols)
+ processed = None
+ self.init_draw(p)
+ gradient = Image.radial_gradient("L").resize(
+ (self.tile_width, self.tile_height), resample=Image.Resampling.BICUBIC)
+ gradient = ImageOps.invert(gradient)
+ p.denoising_strength = self.denoise
+ #p.mask_blur = 0
+ p.mask_blur = self.mask_blur
+
+ for yi in range(rows-1):
+ for xi in range(cols-1):
+ if state.interrupted:
+ break
+ p.width = self.tile_width
+ p.height = self.tile_height
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = 0
+ mask = Image.new("L", (fixed_image.width, fixed_image.height), "black")
+ mask.paste(gradient, (xi*self.tile_width + self.tile_width//2, yi*self.tile_height + self.tile_height//2))
+ p.init_images = [fixed_image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ if len(processed.images) > 0:
+ fixed_image = processed.images[0]
+
+ p.width = fixed_image.width
+ p.height = fixed_image.height
+ if processed is not None:
+ self.initial_info = processed.infotext(p, 0)
+
+ return fixed_image
+
+ def band_pass_process(self, p, image, rows, cols):
+ self.init_draw(p)
+ processed = None
+ p.denoising_strength = self.denoise
+ p.mask_blur = 0
+ gradient = Image.linear_gradient("L")
+ mirror_gradient = Image.new("L", (256, 256), "black")
+ mirror_gradient.paste(gradient.resize((256, 128), resample=Image.Resampling.BICUBIC), (0, 0))
+ mirror_gradient.paste(gradient.rotate(180).resize((256, 128), resample=Image.Resampling.BICUBIC), (0, 128))
+ row_gradient = mirror_gradient.resize((image.width, self.width), resample=Image.Resampling.BICUBIC)
+ col_gradient = mirror_gradient.rotate(90).resize((self.width, image.height), resample=Image.Resampling.BICUBIC)
+
+ for xi in range(1, cols):
+ if state.interrupted:
+ break
+ p.width = self.width + self.padding * 2
+ p.height = image.height
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = self.padding
+ mask = Image.new("L", (image.width, image.height), "black")
+ mask.paste(col_gradient, (xi * self.tile_width - self.width // 2, 0))
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ if len(processed.images) > 0:
+ image = processed.images[0]
+ for yi in range(1, rows):
+ if state.interrupted:
+ break
+ p.width = image.width
+ p.height = self.width + self.padding * 2
+ p.inpaint_full_res = True
+ p.inpaint_full_res_padding = self.padding
+ mask = Image.new("L", (image.width, image.height), "black")
+ mask.paste(row_gradient, (0, yi * self.tile_height - self.width // 2))
+ p.init_images = [image]
+ p.image_mask = mask
+ processed = processing.process_images(p)
+ if len(processed.images) > 0:
+ image = processed.images[0]
+
+ p.width = image.width
+ p.height = image.height
+ if processed is not None:
+ self.initial_info = processed.infotext(p, 0)
+ return image
+
+ def start(self, p, image, rows, cols):
+ if USDUSFMode(self.mode) == USDUSFMode.BAND_PASS:
+ return self.band_pass_process(p, image, rows, cols)
+ elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE:
+ return self.half_tile_process(p, image, rows, cols)
+ elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
+ return self.half_tile_process_corners(p, image, rows, cols)
+ else:
+ return image
+
+class UltimateSDUpscaleScript(scripts_manager.Script):
+ def title(self):
+ return "SD Ultimate Upscale"
+
+ def show(self, is_img2img):
+ return True
+
+ def ui(self, is_img2img): # pylint: disable=unused-argument
+ elem_id_prefix = f"{self.parent}_{elem_id}"
+ target_size_types = ["Default", "Custom", "Rescale"]
+ seams_fix_types = ["None", "Band pass", "Half tile offset pass", "Half tile with intersections"]
+ redraw_modes = ["Linear", "Chess", "None"]
+
+ with gr.Row():
+ gr.HTML('  SD Ultimate Upscale ')
+ with gr.Row():
+ target_size_type = gr.Dropdown(label="Target size", elem_id=f"{elem_id_prefix}_target_size_type", choices=target_size_types, type="index", value=next(iter(target_size_types)))
+ custom_width = gr.Slider(label='Custom width', elem_id=f"{elem_id_prefix}_custom_width", minimum=64, maximum=8192, step=64, value=2048, visible=False, interactive=True)
+ custom_height = gr.Slider(label='Custom height', elem_id=f"{elem_id_prefix}_custom_height", minimum=64, maximum=8192, step=64, value=2048, visible=False, interactive=True)
+ custom_scale = gr.Slider(label='Target scale', elem_id=f"{elem_id_prefix}_custom_scale", minimum=1, maximum=16, step=0.2, value=2, visible=False, interactive=True)
+ with gr.Row():
+ upscaler_index = gr.Dropdown(label='Redraw upscaler', elem_id=f"{elem_id_prefix}_upscaler_index", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")
+ redraw_mode = gr.Dropdown(label="Redraw mode", elem_id=f"{elem_id_prefix}_redraw_mode", choices=redraw_modes, type="index", value=next(iter(redraw_modes)))
+ with gr.Row():
+ tile_width = gr.Slider(elem_id=f"{elem_id_prefix}_tile_width", minimum=0, maximum=2048, step=64, label='Tile width', value=1024)
+ tile_height = gr.Slider(elem_id=f"{elem_id_prefix}_tile_height", minimum=0, maximum=2048, step=64, label='Tile height', value=0)
+ with gr.Row():
+ mask_blur = gr.Slider(elem_id=f"{elem_id_prefix}_mask_blur", label='Tile blur', minimum=0, maximum=64, step=1, value=8)
+ padding = gr.Slider(elem_id=f"{elem_id_prefix}_padding", label='Tile padding', minimum=0, maximum=512, step=1, value=32)
+ with gr.Row():
+ seams_fix_type = gr.Dropdown(label="Seams fix", elem_id=f"{elem_id_prefix}_seams_fix_type", choices=seams_fix_types, type="index", value=next(iter(seams_fix_types)))
+ with gr.Row():
+ seams_fix_denoise = gr.Slider(label='Seams fix denoise', elem_id=f"{elem_id_prefix}_seams_fix_denoise", minimum=0, maximum=1, step=0.01, value=0.35, visible=False, interactive=True)
+ seams_fix_width = gr.Slider(label='Seams fix width', elem_id=f"{elem_id_prefix}_seams_fix_width", minimum=0, maximum=128, step=1, value=64, visible=False, interactive=True)
+ seams_fix_mask_blur = gr.Slider(label='Seams fix mask blur', elem_id=f"{elem_id_prefix}_seams_fix_mask_blur", minimum=0, maximum=64, step=1, value=4, visible=False, interactive=True)
+ seams_fix_padding = gr.Slider(label='Seams fix padding', elem_id=f"{elem_id_prefix}_seams_fix_padding", minimum=0, maximum=128, step=1, value=16, visible=False, interactive=True)
+
+ def select_fix_type(fix_index):
+ all_visible = fix_index != 0
+ mask_blur_visible = fix_index == 2 or fix_index == 3
+ width_visible = fix_index == 1
+ return [gr.update(visible=all_visible), gr.update(visible=width_visible), gr.update(visible=mask_blur_visible), gr.update(visible=all_visible)]
+
+ seams_fix_type.change(fn=select_fix_type, inputs=seams_fix_type, outputs=[seams_fix_denoise, seams_fix_width, seams_fix_mask_blur, seams_fix_padding])
+
+ def select_scale_type(scale_index):
+ is_custom_size = scale_index == 1
+ is_custom_scale = scale_index == 2
+ return [gr.update(visible=is_custom_size), gr.update(visible=is_custom_size), gr.update(visible=is_custom_scale)]
+
+ target_size_type.change(fn=select_scale_type, inputs=target_size_type, outputs=[custom_width, custom_height, custom_scale])
+
+ def init_field(scale_name):
+ try:
+ scale_index = target_size_types.index(scale_name)
+ custom_width.visible = custom_height.visible = scale_index == 1
+ custom_scale.visible = scale_index == 2
+ except Exception:
+ pass
+
+ target_size_type.init_field = init_field
+
+ return [tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
+ upscaler_index, redraw_mode, seams_fix_mask_blur,
+ seams_fix_type, target_size_type, custom_width, custom_height, custom_scale]
+
+ def run(self, p: StableDiffusionProcessing, *args):
+ tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding, upscaler_index, redraw_mode, seams_fix_mask_blur, seams_fix_type, target_size_type, custom_width, custom_height, custom_scale = args
+
+ if p.init_images is None or len(p.init_images) == 0: # Dont run on t2i
+ return None
+
+ # Init
+ processing.fix_seed(p)
+
+ p.do_not_save_grid = True
+ p.do_not_save_samples = True
+ p.inpaint_full_res = False
+ p.inpainting_fill = 1
+ p.n_iter = 1
+ p.batch_size = 1
+
+ # Init image
+ init_img = p.init_images[0]
+ if init_img is None:
+ return Processed(p, [], p.seed, "Empty image")
+ init_img = images.flatten(init_img, opts.img2img_background_color)
+
+ #override size
+ if target_size_type == 1:
+ p.width = custom_width
+ p.height = custom_height
+ if target_size_type == 2:
+ p.width = math.ceil((init_img.width * custom_scale) / 64) * 64
+ p.height = math.ceil((init_img.height * custom_scale) / 64) * 64
+
+ # Upscaling
+ upscaler = USDUpscaler(p, init_img, upscaler_index, tile_width, tile_height)
+ upscaler.upscale()
+
+ # Drawing
+ upscaler.setup_redraw(redraw_mode, padding, mask_blur)
+ upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type)
+ upscaler.add_extra_info()
+ upscaler.process()
+ result_images = upscaler.result_images
+
+ return Processed(p, result_images, p.seed, upscaler.initial_info if upscaler.initial_info is not None else "")
diff --git a/scripts/xyz/xyz_grid_draw.py b/scripts/xyz/xyz_grid_draw.py
index 45ab14cbb..10477fc7c 100644
--- a/scripts/xyz/xyz_grid_draw.py
+++ b/scripts/xyz/xyz_grid_draw.py
@@ -7,7 +7,7 @@ from modules.logger import log
from modules.image.util import draw_text
-def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid: False, include_time: False, include_text: False): # pylint: disable=unused-argument
+def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid=False, include_time=False, include_text=False): # pylint: disable=unused-argument
x_texts = [[GridAnnotation(x)] for x in x_labels]
y_texts = [[GridAnnotation(y)] for y in y_labels]
z_texts = [[GridAnnotation(z)] for z in z_labels]
diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py
new file mode 100644
index 000000000..c84546943
--- /dev/null
+++ b/test/test-flux2-native-adapters.py
@@ -0,0 +1,1268 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for Flux2/Klein native adapter loaders.
+
+Covers the nine native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3,
+GLoRA, Norm, Full) plus DoRA threading via the universal
+NetworkModule.finalize_updown hook, and ex_bias accumulation across
+stacked Norm adapters.
+
+The tests build a mock Flux2-shaped transformer, write synthetic
+safetensors files for each adapter format, and exercise the full loader
+path: state-dict -> group_by_suffixes -> resolve_targets ->
+NetworkModule* -> calc_updown -> finalize_updown.
+
+Loader correctness is verified against the documented LyCORIS / kohya
+save formats. End-to-end inference verification against real adapter
+files remains an open gap for the rarer families.
+
+No running server required.
+
+Usage:
+ python test/test-flux2-native-adapters.py
+"""
+
+import os
+import sys
+import tempfile
+import time
+
+import torch
+import safetensors.torch
+
+script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, script_dir)
+os.chdir(script_dir)
+
+os.environ['SD_INSTALL_QUIET'] = '1'
+
+# Bootstrap cmd_args before any module that pulls in shared.py.
+# parse_args() registers the full main+compat option set; installer.add_args
+# adds the installer-specific options on top.
+import modules.cmd_args # pylint: disable=wrong-import-position
+import installer # pylint: disable=wrong-import-position
+_orig_argv = sys.argv
+sys.argv = [sys.argv[0]]
+try:
+ modules.cmd_args.parse_args()
+finally:
+ sys.argv = _orig_argv
+installer.add_args(modules.cmd_args.parser)
+modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
+
+from modules.errors import log # pylint: disable=wrong-import-position
+from modules import shared # pylint: disable=wrong-import-position
+from modules.lora import ( # pylint: disable=wrong-import-position
+ network, network_lora, network_lokr, network_hada, network_oft, network_boft,
+ network_ia3, network_glora, network_norm, network_full,
+)
+from modules.lora import lora_apply # pylint: disable=wrong-import-position
+from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position
+from pipelines.flux import flux2_lora as F # pylint: disable=wrong-import-position
+
+
+# ============================================================
+# Test infrastructure
+# ============================================================
+
+results: dict[str, dict] = {}
+
+
+def category(name: str):
+ if name not in results:
+ results[name] = {'passed': 0, 'failed': 0, 'tests': []}
+ return name
+
+
+def record(cat: str, passed: bool, name: str, detail: str = ''):
+ status = 'PASS' if passed else 'FAIL'
+ results[cat]['passed' if passed else 'failed'] += 1
+ results[cat]['tests'].append((status, name))
+ msg = f' {status}: {name}'
+ if detail:
+ msg += f' ({detail})'
+ if passed:
+ log.info(msg)
+ else:
+ log.error(msg)
+
+
+def run_test(cat: str, fn):
+ name = fn.__name__
+ try:
+ ok = fn()
+ if ok is False:
+ record(cat, False, name)
+ else:
+ record(cat, True, name)
+ except AssertionError as e:
+ record(cat, False, name, str(e))
+ except Exception as e: # pylint: disable=broad-except
+ record(cat, False, name, f'exception: {e}')
+ import traceback
+ traceback.print_exc()
+
+
+# ============================================================
+# Mock Flux2-shaped transformer
+# ============================================================
+
+# Small but realistic shapes so tests are fast.
+HIDDEN = 128
+QKV_OUT = 128
+MLP_OUT = 512
+HEAD_DIM = 32
+SINGLE_FUSED_OUT = 3 * QKV_OUT + MLP_OUT # to_qkv_mlp_proj
+SINGLE_OUT_IN = QKV_OUT + MLP_OUT # to_out
+N_DOUBLE = 2
+N_SINGLE = 2
+
+
+# pylint: disable=attribute-defined-outside-init
+# `_Holder` is a generic container we attach diffusers-shaped children to at
+# runtime to mirror Flux2's transformer module tree.
+class _Holder(torch.nn.Module):
+ """Empty container module — we attach Linear/RMSNorm children dynamically."""
+
+
+def build_mock_transformer():
+ """Build a torch.nn.Module mimicking Flux2's diffusers-side module tree.
+
+ Mirrors the paths in F2_SINGLE_MAP / F2_DOUBLE_MAP / F2_QKV_MAP plus the
+ RMSNorm targets inside attention so that try_load_norm has something to
+ bind. Sizes are scaled-down but proportional so chunking math is realistic.
+ """
+ transformer = _Holder()
+ transformer.transformer_blocks = torch.nn.ModuleList()
+ for _ in range(N_DOUBLE):
+ block = _Holder()
+ block.attn = _Holder()
+ block.attn.to_q = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.to_k = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.to_v = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.to_out = torch.nn.ModuleList([torch.nn.Linear(QKV_OUT, HIDDEN, bias=False)])
+ block.attn.add_q_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.add_k_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.add_v_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
+ block.attn.to_add_out = torch.nn.Linear(QKV_OUT, HIDDEN, bias=False)
+ block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM)
+ block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM)
+ block.attn.norm_added_q = torch.nn.RMSNorm(HEAD_DIM)
+ block.attn.norm_added_k = torch.nn.RMSNorm(HEAD_DIM)
+ block.ff = _Holder()
+ block.ff.linear_in = torch.nn.Linear(HIDDEN, MLP_OUT, bias=False)
+ block.ff.linear_out = torch.nn.Linear(MLP_OUT, HIDDEN, bias=False)
+ block.ff_context = _Holder()
+ block.ff_context.linear_in = torch.nn.Linear(HIDDEN, MLP_OUT, bias=False)
+ block.ff_context.linear_out = torch.nn.Linear(MLP_OUT, HIDDEN, bias=False)
+ transformer.transformer_blocks.append(block)
+ transformer.single_transformer_blocks = torch.nn.ModuleList()
+ for _ in range(N_SINGLE):
+ sblock = _Holder()
+ sblock.attn = _Holder()
+ sblock.attn.to_qkv_mlp_proj = torch.nn.Linear(HIDDEN, SINGLE_FUSED_OUT, bias=False)
+ sblock.attn.to_out = torch.nn.Linear(SINGLE_OUT_IN, HIDDEN, bias=False)
+ transformer.single_transformer_blocks.append(sblock)
+ return transformer
+
+
+class _MockFlux2Pipeline:
+ """Class name carries 'Flux2' so shared.sd_model_type returns 'f2'
+ via modeldata.get_model_type's name-based dispatch."""
+
+ def __init__(self, transformer):
+ self.transformer = transformer
+ self.text_encoder = None
+
+
+class _MockFlux2SdModel:
+ """Outer wrapper holding the pipe + the network_layer_mapping that
+ lora_convert.assign_network_names_to_compvis_modules writes onto."""
+
+ def __init__(self, pipe):
+ self.pipe = pipe
+ self.network_layer_mapping = {}
+ self.embedding_db = None
+ self.__class__.__name__ = 'Flux2Pipeline' # belt-and-suspenders
+
+
+def install_mock_pipe():
+ """Set shared.sd_model to a mock exposing a Flux2-shaped transformer.
+
+ Each test calls this fresh so module.network_layer_name stamps from prior
+ tests don't leak (notably the loader-local stamping in try_load_norm).
+
+ Bypasses the ModelData lock by writing directly to model_data.sd_model;
+ shared.sd_model = ... goes through set_sd_model which is a no-op when
+ model_data.locked is True (the default outside webui startup).
+ """
+ transformer = build_mock_transformer()
+ pipe = _MockFlux2Pipeline(transformer)
+ sd_model = _MockFlux2SdModel(pipe)
+ from modules.modeldata import model_data
+ model_data.sd_model = sd_model
+ return sd_model
+
+
+# ============================================================
+# State-dict synthesizers (one per family)
+# ============================================================
+
+# Realistic ranks for the synthesized factors.
+RANK_LORA = 8
+RANK_LOKR = 4
+LOKR_W1_DIM = 8
+LOKR_W2_DIM = HIDDEN // LOKR_W1_DIM
+
+
+def sd_lora_kohya_qkv():
+ """Kohya-format LoRA targeting fused img_attn.qkv (3-way chunked at load)."""
+ return {
+ # Down weight is shared across Q/K/V; up weight is fused (3*QKV_OUT, RANK).
+ 'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(3 * QKV_OUT, RANK_LORA),
+ 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_lora_bfl_proj():
+ """BFL/AI-toolkit-format LoRA on a non-fused double_blocks proj target."""
+ return {
+ 'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, QKV_OUT),
+ 'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
+ 'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_lora_peft_to_q():
+ """Diffusers PEFT-format LoRA targeting a single split-QKV diffusers path."""
+ return {
+ 'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
+ }
+
+
+def sd_lora_bare_bfl_mlp():
+ """Bare BFL keys (no prefix) on an MLP target — exercises the bare fallback."""
+ return {
+ 'double_blocks.0.img_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'double_blocks.0.img_mlp.0.lora_B.weight': torch.randn(MLP_OUT, RANK_LORA),
+ }
+
+
+def sd_lora_peft_saved_fal_style():
+ """PEFT-saved LoRA: base_model.model. wrapper around BFL paths.
+
+ Mirrors the actual key layout of fal/flux-2-klein-4B-outpaint-lora as
+ inspected from its safetensors header on HF — first 12 keys all of the
+ form ``base_model.model.double_blocks.0.img_attn.proj.lora_A.weight``.
+ """
+ return {
+ 'base_model.model.double_blocks.1.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, QKV_OUT),
+ 'base_model.model.double_blocks.1.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
+ 'base_model.model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_lora_peft_saved_dreambooth_style():
+ """PEFT-saved LoRA where the wrapped path is diffusers-form (HF DreamBooth).
+
+ HF DreamBooth scripts target diffusers modules by name (e.g.
+ ``transformer.transformer_blocks.0.attn.to_q``); after peft.save_pretrained
+ wraps the in-memory keys, the saved form is
+ ``base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight``.
+ """
+ return {
+ 'base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
+ }
+
+
+def sd_lora_peft_saved_diffusers_style():
+ """PEFT-saved LoRA with bare-diffusers paths and a ``.default.`` adapter-name infix.
+
+ Mirrors the actual key layout of community Flux2 LoRAs exported via
+ ``Flux2Transformer2DModel.save_lora_adapter()`` after attaching an unnamed
+ PEFT adapter (e.g. ``lenovo_flux_klein9b.safetensors``):
+
+ - No ``transformer.`` / ``diffusion_model.`` / ``base_model.model.`` wrapper
+ - Keys start with bare ``transformer_blocks.`` or ``single_transformer_blocks.``
+ - ``.lora_A.default.weight`` / ``.lora_B.default.weight`` (PEFT adapter slot)
+ """
+ return {
+ # Single-block fused module — exercises the bare-diffusers prefix branch.
+ 'single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_A.default.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_B.default.weight': torch.randn(SINGLE_FUSED_OUT, RANK_LORA),
+ # Double-block split target — exercises both the bare-diffusers branch and the .default. strip.
+ 'transformer_blocks.0.attn.to_q.lora_A.default.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'transformer_blocks.0.attn.to_q.lora_B.default.weight': torch.randn(QKV_OUT, RANK_LORA),
+ }
+
+
+def sd_lora_with_dora_scale():
+ """LoRA with a dora_scale companion vector — exercises DoRA threading."""
+ return {
+ 'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
+ 'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
+ 'transformer.transformer_blocks.0.attn.to_q.dora_scale': torch.randn(QKV_OUT),
+ }
+
+
+def sd_lokr_bfl_proj():
+ """BFL-format LoKR on a non-fused proj target."""
+ return {
+ # w1 = (LOKR_W1_DIM, LOKR_W1_DIM), w2 = (LOKR_W2_DIM, LOKR_W2_DIM); kron is (HIDDEN, HIDDEN).
+ # Target shape is (HIDDEN, QKV_OUT) here; we need shape compat for calc_updown.
+ # Use w1=(QKV_OUT/LOKR_W2_DIM, HIDDEN/LOKR_W2_DIM) outer factor; LyCORIS picks dims by factorization.
+ 'diffusion_model.double_blocks.1.img_attn.proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
+ 'diffusion_model.double_blocks.1.img_attn.proj.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
+ 'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(LOKR_W1_DIM)),
+ }
+
+
+def sd_lokr_kohya_qkv():
+ """Kohya LoKR targeting fused img_attn.qkv — exercises NetworkModuleLokrChunk."""
+ # Target after kron is (3*QKV_OUT, HIDDEN); chunk along dim=0 gives Q/K/V slices.
+ fused_out = 3 * QKV_OUT
+ # Pick w1=(8, ...) so kron of w1 with w2 spans fused_out x HIDDEN.
+ w1_dim_out = 8
+ w2_dim_out = fused_out // w1_dim_out
+ w1_dim_in = 4
+ w2_dim_in = HIDDEN // w1_dim_in
+ return {
+ 'lora_unet_double_blocks_0_img_attn_qkv.lokr_w1': torch.randn(w1_dim_out, w1_dim_in),
+ 'lora_unet_double_blocks_0_img_attn_qkv.lokr_w2': torch.randn(w2_dim_out, w2_dim_in),
+ 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(w1_dim_in)),
+ }
+
+
+def sd_lokr_simpletuner_lycoris_style():
+ """LyCORIS-standalone-format LoKR (SimpleTuner save).
+
+ Mirrors the actual key layout of markury/flux2k9b-simpletuner-lokr-loona
+ as inspected from its safetensors header on HF — keys of the form
+ ``lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1`` (288 such keys
+ in the real file). The path under ``lycoris_`` is a diffusers path with
+ dots rendered as underscores; resolve_targets returns it verbatim and
+ the caller's ``.replace('.', '_')`` is a no-op.
+
+ Two targets here exercise both an img-side projection (to_q) and a
+ txt-side projection (add_k_proj) which uses underscores in the module
+ name itself — the latter is the case where naïve underscore-to-dot
+ expansion would corrupt the path.
+ """
+ return {
+ 'lycoris_transformer_blocks_0_attn_to_q.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
+ 'lycoris_transformer_blocks_0_attn_to_q.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
+ 'lycoris_transformer_blocks_0_attn_to_q.alpha': torch.tensor(float(LOKR_W1_DIM)),
+ 'lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
+ 'lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
+ 'lycoris_transformer_blocks_0_attn_add_k_proj.alpha': torch.tensor(float(LOKR_W1_DIM)),
+ }
+
+
+def sd_loha_bfl_proj():
+ """BFL-format LoHA on a non-fused proj target."""
+ return {
+ 'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_a': torch.randn(HIDDEN, RANK_LORA),
+ 'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_b': torch.randn(RANK_LORA, QKV_OUT),
+ 'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_a': torch.randn(HIDDEN, RANK_LORA),
+ 'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_b': torch.randn(RANK_LORA, QKV_OUT),
+ 'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_loha_kohya_qkv():
+ """Kohya LoHA targeting fused img_attn.qkv — exercises NetworkModuleHadaChunk."""
+ fused_out = 3 * QKV_OUT
+ return {
+ 'lora_unet_double_blocks_0_img_attn_qkv.hada_w1_a': torch.randn(fused_out, RANK_LORA),
+ 'lora_unet_double_blocks_0_img_attn_qkv.hada_w1_b': torch.randn(RANK_LORA, HIDDEN),
+ 'lora_unet_double_blocks_0_img_attn_qkv.hada_w2_a': torch.randn(fused_out, RANK_LORA),
+ 'lora_unet_double_blocks_0_img_attn_qkv.hada_w2_b': torch.randn(RANK_LORA, HIDDEN),
+ 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_oft_kohya_proj():
+ """Kohya-format OFT (oft_blocks) on a non-fused proj target."""
+ # OFT block-diagonal: out_dim = HIDDEN, num_blocks=8, block_size=16.
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.oft_blocks': torch.zeros(8, 16, 16),
+ 'lora_unet_double_blocks_1_img_attn_proj.alpha': torch.tensor(1e-3),
+ }
+
+
+def sd_oft_lycoris_proj():
+ """LyCORIS-format OFT (oft_diag) — exercises the network_oft.py:58 NPE-fix path.
+
+ LyCORIS reads ``self.dim = oft_diag.shape[1]`` (the block_size) and computes
+ ``block_size, num_blocks = factorization(out_dim, dim)``. ``factorization``
+ always returns ``(smaller, larger)``. For out_dim=HIDDEN=128 with dim=8, it
+ returns (8, 16) ⇒ block_size=8, num_blocks=16. The tensor must therefore
+ have shape ``(num_blocks, block_size, block_size) = (16, 8, 8)``.
+ """
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.oft_diag': torch.zeros(16, 8, 8),
+ }
+
+
+def sd_oft_kohya_qkv_skipped():
+ """OFT on fused img_attn.qkv — must be skipped with a warning (no chunk class)."""
+ return {
+ 'lora_unet_double_blocks_0_img_attn_qkv.oft_blocks': torch.zeros(8, 48, 48),
+ 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(1e-3),
+ }
+
+
+def sd_boft_butterfly():
+ """BOFT (butterfly-OFT) — re-uses ``oft_blocks`` key with 4-D shape.
+
+ Mirrors the LyCORIS upstream save layout (boft.py weight_list:
+ ``oft_blocks, rescale, alpha``; tensor shape per __init__:
+ ``(boft_m, block_num, block_size, block_size)``).
+
+ For HIDDEN=128 with block_size=16, block_num=8:
+ - boft_m derives from sum-of-set-bits of (block_num-1) + 1 = bits(7)+1 = 4
+ - cascade indices: k = 2^i * r_b for i in 0..3, with r_b = 8
+ - all four stages divide cleanly into 128-row tensors
+ """
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.oft_blocks':
+ torch.randn(4, 8, 16, 16) * 0.01,
+ 'lora_unet_double_blocks_1_img_attn_proj.alpha':
+ torch.tensor(1e-3),
+ }
+
+
+def sd_ia3_proj():
+ """IA3 on a non-fused target. on_input=False ⇒ output-axis scaling."""
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.weight': torch.randn(HIDDEN),
+ 'lora_unet_double_blocks_1_img_attn_proj.on_input': torch.tensor(0),
+ }
+
+
+def sd_ia3_qkv_skipped():
+ """IA3 on fused QKV — must be skipped with a warning."""
+ return {
+ 'lora_unet_double_blocks_0_img_attn_qkv.weight': torch.randn(3 * QKV_OUT),
+ 'lora_unet_double_blocks_0_img_attn_qkv.on_input': torch.tensor(0),
+ }
+
+
+def sd_glora_proj():
+ """GLoRA on a non-fused proj target.
+
+ Target host module is ``attn.to_out.0`` with weight shape (out, in)
+ = (HIDDEN, QKV_OUT). NetworkModuleGLora computes
+ ``updown = (w2b @ w1b) + ((target @ w2a) @ w1a)``, which needs:
+
+ - w2b: (out, rank), w1b: (rank, in) ⇒ first term shape (out, in)
+ - w2a: (in, rank), w1a: (rank, in) ⇒ second term ((out,in) @ (in,rank)) @ (rank,in) = (out, in)
+ """
+ out_dim = HIDDEN
+ in_dim = QKV_OUT
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.a1.weight': torch.randn(RANK_LORA, in_dim),
+ 'lora_unet_double_blocks_1_img_attn_proj.a2.weight': torch.randn(in_dim, RANK_LORA),
+ 'lora_unet_double_blocks_1_img_attn_proj.b1.weight': torch.randn(RANK_LORA, in_dim),
+ 'lora_unet_double_blocks_1_img_attn_proj.b2.weight': torch.randn(out_dim, RANK_LORA),
+ 'lora_unet_double_blocks_1_img_attn_proj.alpha': torch.tensor(float(RANK_LORA)),
+ }
+
+
+def sd_norm_peft_attn_norm_q():
+ """Norm targeting a real Flux2 RMSNorm (attn.norm_q)."""
+ return {
+ 'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.randn(HEAD_DIM) * 0.01,
+ 'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.randn(HEAD_DIM) * 0.01,
+ }
+
+
+def sd_full_proj():
+ """Full-rank delta on a proj target. diff matches host weight shape (out, in)."""
+ return {
+ 'lora_unet_double_blocks_1_img_attn_proj.diff': torch.randn(HIDDEN, QKV_OUT) * 0.01,
+ 'lora_unet_double_blocks_1_img_attn_proj.diff_b': torch.randn(HIDDEN) * 0.01,
+ }
+
+
+# ============================================================
+# Test plumbing helpers
+# ============================================================
+
+class TempLora:
+ """Context-manager that materializes a synthetic state dict as a real
+ safetensors file and yields a NetworkOnDisk pointing at it.
+
+ Required because flux2_lora's loaders read via sd_models.read_state_dict,
+ which only accepts a filename. Cleans up the temp file on exit.
+ """
+
+ def __init__(self, state_dict, name='test'):
+ self.state_dict = state_dict
+ self.name = name
+ self.path = None
+
+ def __enter__(self):
+ fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_')
+ os.close(fd)
+ # safetensors requires contiguous tensors.
+ sd = {k: v.contiguous() for k, v in self.state_dict.items()}
+ safetensors.torch.save_file(sd, self.path)
+ return _MockNetworkOnDisk(self.name, self.path)
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ if self.path and os.path.exists(self.path):
+ os.unlink(self.path)
+
+
+class _MockNetworkOnDisk:
+ """Minimal NetworkOnDisk shim — loaders only touch .filename, .name, .shorthash, .sd_version."""
+
+ def __init__(self, name, filename):
+ self.name = name
+ self.filename = filename
+ self.shorthash = ''
+ self.sd_version = 'f2'
+ self.metadata = {}
+
+ def read_hash(self):
+ pass
+
+
+def assert_shape(t: torch.Tensor, expected_shape, label=''):
+ actual = tuple(t.shape)
+ expected = tuple(expected_shape)
+ assert actual == expected, f'{label} shape {actual} != expected {expected}'
+
+
+def assert_finite(t: torch.Tensor, label=''):
+ assert torch.isfinite(t).all(), f'{label} contains non-finite values'
+
+
+def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0):
+ """Wire a NetworkModule into a Network with given multipliers so calc_updown
+ has a meaningful multiplier()/calc_scale() context."""
+ net_module.network.te_multiplier = te_mul
+ net_module.network.unet_multiplier = unet_mul
+ return net_module
+
+
+# ============================================================
+# Tests — parsing primitives
+# ============================================================
+
+CAT_PARSE = category('parse')
+
+
+def test_parse_key_all_prefixes():
+ cases = [
+ ('lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight',
+ F.LORA_SUFFIXES,
+ ('lora_unet_', 'double_blocks_0_img_attn_qkv', 'lora_down.weight')),
+ ('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight',
+ F.LORA_SUFFIXES,
+ ('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
+ ('transformer.transformer_blocks.0.attn.to_q.lora_B.weight',
+ F.LORA_SUFFIXES,
+ ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
+ ('double_blocks.10.img_mlp.0.lora_A.weight',
+ F.LORA_SUFFIXES,
+ (None, 'double_blocks.10.img_mlp.0', 'lora_down.weight')),
+ ('random.unrelated.key', F.LORA_SUFFIXES, None),
+ ]
+ for key, suffixes, expected in cases:
+ got = F.parse_key(key, suffixes)
+ assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
+ return True
+
+
+def test_resolve_targets_qkv_chunking():
+ # Kohya double_blocks fused QKV → three chunks targeting Q/K/V.
+ targets = F.resolve_targets('lora_unet_', 'double_blocks_0_img_attn_qkv')
+ assert targets == [
+ ('transformer_blocks.0.attn.to_q', 0, 3),
+ ('transformer_blocks.0.attn.to_k', 1, 3),
+ ('transformer_blocks.0.attn.to_v', 2, 3),
+ ], f'kohya img_attn.qkv → {targets}'
+
+ targets = F.resolve_targets('lora_unet_', 'double_blocks_5_txt_attn_qkv')
+ assert targets == [
+ ('transformer_blocks.5.attn.add_q_proj', 0, 3),
+ ('transformer_blocks.5.attn.add_k_proj', 1, 3),
+ ('transformer_blocks.5.attn.add_v_proj', 2, 3),
+ ], f'kohya txt_attn.qkv → {targets}'
+
+ targets = F.resolve_targets('diffusion_model.', 'single_blocks.7.linear1')
+ assert targets == [('single_transformer_blocks.7.attn.to_qkv_mlp_proj', None, None)]
+
+ targets = F.resolve_targets('transformer.', 'transformer_blocks.0.attn.to_q')
+ assert targets == [('transformer_blocks.0.attn.to_q', None, None)]
+
+ targets = F.resolve_targets('weird_prefix.', 'whatever')
+ assert targets == []
+ return True
+
+
+def test_parse_key_peft_wrapper_unwrap():
+ """base_model.model. wrapper is stripped before format detection.
+
+ Verifies the unwrap handles all four content cases:
+ - BFL keys under the wrapper (fal style)
+ - Diffusers paths under transformer. under the wrapper (HF DreamBooth style)
+ - Bare BFL under the wrapper
+ - Wrapper not present (passthrough)
+ """
+ cases = [
+ ('base_model.model.double_blocks.1.img_attn.proj.lora_A.weight',
+ F.LORA_SUFFIXES,
+ (None, 'double_blocks.1.img_attn.proj', 'lora_down.weight')),
+ ('base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight',
+ F.LORA_SUFFIXES,
+ ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
+ ('base_model.model.lora_unet_double_blocks_0_img_attn_proj.lora_down.weight',
+ F.LORA_SUFFIXES,
+ ('lora_unet_', 'double_blocks_0_img_attn_proj', 'lora_down.weight')),
+ ('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight', # no wrapper
+ F.LORA_SUFFIXES,
+ ('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
+ ]
+ for key, suffixes, expected in cases:
+ got = F.parse_key(key, suffixes)
+ assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
+ return True
+
+
+def test_parse_key_lycoris_prefix():
+ """lycoris_ prefix yields an underscored diffusers path; no _ -> . conversion."""
+ cases = [
+ ('lycoris_transformer_blocks_0_attn_to_q.lokr_w1',
+ F.LOKR_SUFFIXES,
+ ('lycoris_', 'transformer_blocks_0_attn_to_q', 'lokr_w1')),
+ ('lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2',
+ F.LOKR_SUFFIXES,
+ ('lycoris_', 'transformer_blocks_0_attn_add_k_proj', 'lokr_w2')),
+ ('lycoris_single_transformer_blocks_5_attn_to_qkv_mlp_proj.lokr_w1',
+ F.LOKR_SUFFIXES,
+ ('lycoris_', 'single_transformer_blocks_5_attn_to_qkv_mlp_proj', 'lokr_w1')),
+ ]
+ for key, suffixes, expected in cases:
+ got = F.parse_key(key, suffixes)
+ assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
+
+ # resolve_targets: the underscored path is returned verbatim (no chunk).
+ targets = F.resolve_targets('lycoris_', 'transformer_blocks_0_attn_add_k_proj')
+ assert targets == [('transformer_blocks_0_attn_add_k_proj', None, None)], f'targets={targets}'
+ return True
+
+
+def test_parse_key_bare_diffusers_and_peft_default():
+ """Bare-diffusers paths + ``.lora_[AB]..weight`` infix.
+
+ Covers ``Flux2Transformer2DModel.save_lora_adapter()`` output where the file
+ has no wrapping prefix and PEFT keeps the adapter slot name in the suffix.
+ """
+ bd = F.BARE_DIFFUSERS_PREFIX_USED
+ cases = [
+ ('single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_A.default.weight',
+ F.LORA_SUFFIXES,
+ (bd, 'single_transformer_blocks.0.attn.to_qkv_mlp_proj', 'lora_down.weight')),
+ ('transformer_blocks.0.attn.to_q.lora_B.default.weight',
+ F.LORA_SUFFIXES,
+ (bd, 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
+ # Non-default adapter name also strips.
+ ('transformer_blocks.7.attn.to_v.lora_A.style.weight',
+ F.LORA_SUFFIXES,
+ (bd, 'transformer_blocks.7.attn.to_v', 'lora_down.weight')),
+ # No adapter-name infix passes through unchanged.
+ ('transformer_blocks.0.attn.to_q.lora_A.weight',
+ F.LORA_SUFFIXES,
+ (bd, 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
+ # Wrapped path with adapter-name infix: unwrap + strip both apply.
+ ('base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.default.weight',
+ F.LORA_SUFFIXES,
+ ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
+ ]
+ for key, suffixes, expected in cases:
+ got = F.parse_key(key, suffixes)
+ assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
+
+ # resolve_targets passes the bare-diffusers path through verbatim.
+ targets = F.resolve_targets(bd, 'single_transformer_blocks.5.attn.to_out')
+ assert targets == [('single_transformer_blocks.5.attn.to_out', None, None)], f'targets={targets}'
+ return True
+
+
+def test_marker_disambiguation():
+ # Each family's marker must reject other families' files.
+ pure_lora = {'lora_unet_x.lora_down.weight': torch.zeros(1, 1),
+ 'lora_unet_x.lora_up.weight': torch.zeros(1, 1)}
+ assert F.has_marker(pure_lora, F.LORA_MARKERS)
+ assert not F.has_marker(pure_lora, F.LOKR_MARKERS)
+ assert not F.has_marker(pure_lora, F.LOHA_MARKERS)
+ assert not F.has_marker(pure_lora, F.OFT_MARKERS)
+ assert not F.has_marker(pure_lora, F.IA3_MARKERS)
+ assert not F.has_marker(pure_lora, F.GLORA_MARKERS)
+ assert not F.has_marker(pure_lora, F.NORM_MARKERS)
+ assert not F.has_marker(pure_lora, F.FULL_MARKERS)
+
+ ia3 = {'lora_unet_x.weight': torch.zeros(1), 'lora_unet_x.on_input': torch.tensor(1)}
+ assert F.has_marker(ia3, F.IA3_MARKERS)
+ # IA3 shares ".weight" but only on_input is the disambiguator.
+ return True
+
+
+# ============================================================
+# Tests — loaders (one per family)
+# ============================================================
+
+CAT_LOADER = category('loader')
+
+
+def _load_via(try_fn, state_dict, name='test'):
+ install_mock_pipe()
+ with TempLora(state_dict, name=name) as nod:
+ return try_fn(name, nod, lora_scale=1.0)
+
+
+def test_lora_kohya_fused_qkv_chunked():
+ net = _load_via(F.try_load_lora, sd_lora_kohya_qkv())
+ assert net is not None and len(net.modules) == 3, f'expected 3 chunked modules, got {net.modules if net else None}'
+ # Network keys correspond to the three split projections.
+ expected_keys = {
+ 'lora_transformer_transformer_blocks_0_attn_to_q',
+ 'lora_transformer_transformer_blocks_0_attn_to_k',
+ 'lora_transformer_transformer_blocks_0_attn_to_v',
+ }
+ assert set(net.modules) == expected_keys, f'got {set(net.modules)}'
+ # Every module is a NetworkModuleLora; up-weight has been chunked along dim 0.
+ for nk, mod in net.modules.items():
+ assert isinstance(mod, network_lora.NetworkModuleLora), f'{nk}: type={type(mod).__name__}'
+ # The chunked up tensor has shape (QKV_OUT, RANK), not (3*QKV_OUT, RANK).
+ up_shape = tuple(mod.up_model.weight.shape)
+ assert up_shape == (QKV_OUT, RANK_LORA), f'{nk}: up shape {up_shape}'
+ return True
+
+
+def test_lora_bfl_non_fused():
+ net = _load_via(F.try_load_lora, sd_lora_bfl_proj())
+ assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
+ nk, mod = next(iter(net.modules.items()))
+ assert nk == 'lora_transformer_transformer_blocks_1_attn_to_out_0', f'{nk}'
+ assert isinstance(mod, network_lora.NetworkModuleLora)
+ return True
+
+
+def test_lora_peft_format():
+ net = _load_via(F.try_load_lora, sd_lora_peft_to_q())
+ assert net is not None and len(net.modules) == 1
+ assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
+ return True
+
+
+def test_lora_bare_bfl_format():
+ net = _load_via(F.try_load_lora, sd_lora_bare_bfl_mlp())
+ assert net is not None and len(net.modules) == 1
+ assert 'lora_transformer_transformer_blocks_0_ff_linear_in' in net.modules
+ return True
+
+
+def test_lora_peft_saved_fal_style():
+ """Real-world PEFT-saved LoRA with BFL paths under base_model.model. wrapper."""
+ net = _load_via(F.try_load_lora, sd_lora_peft_saved_fal_style())
+ assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
+ # base_model.model.double_blocks.1.img_attn.proj → transformer_blocks.1.attn.to_out.0
+ assert 'lora_transformer_transformer_blocks_1_attn_to_out_0' in net.modules
+ return True
+
+
+def test_lora_peft_saved_dreambooth_style():
+ """PEFT-saved LoRA with diffusers paths under base_model.model. wrapper.
+
+ HF DreamBooth-style: the path under base_model.model. starts with
+ transformer., which the unwrap+reparse routes through the standard PEFT
+ branch. Verifies the unwrap handles nested format prefixes correctly.
+ """
+ net = _load_via(F.try_load_lora, sd_lora_peft_saved_dreambooth_style())
+ assert net is not None and len(net.modules) == 1
+ assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
+ return True
+
+
+def test_lora_peft_saved_diffusers_style():
+ """Bare-diffusers paths + PEFT ``.default.`` adapter-name infix.
+
+ Exercises the bare-diffusers prefix branch (``single_transformer_blocks.``,
+ ``transformer_blocks.``) combined with the PEFT-saved ``.lora_A.default.weight``
+ suffix shape produced by ``Flux2Transformer2DModel.save_lora_adapter()``.
+ """
+ net = _load_via(F.try_load_lora, sd_lora_peft_saved_diffusers_style())
+ assert net is not None and len(net.modules) == 2, f'expected 2 modules, got {net.modules if net else None}'
+ assert 'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj' in net.modules
+ assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
+ return True
+
+
+def test_lora_dora_threading():
+ net = _load_via(F.try_load_lora, sd_lora_with_dora_scale())
+ assert net is not None and len(net.modules) == 1
+ mod = next(iter(net.modules.values()))
+ assert mod.dora_scale is not None, 'dora_scale not threaded into NetworkModule'
+ # Apply path runs through finalize_updown which calls apply_weight_decompose.
+ target = torch.randn(QKV_OUT, HIDDEN)
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, (QKV_OUT, HIDDEN), 'DoRA updown')
+ assert_finite(updown, 'DoRA updown')
+ return True
+
+
+def test_lokr_bfl_non_fused():
+ net = _load_via(F.try_load_lokr, sd_lokr_bfl_proj())
+ assert net is not None and len(net.modules) == 1
+ nk, mod = next(iter(net.modules.items()))
+ assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk), \
+ f'{nk}: type={type(mod).__name__}'
+ return True
+
+
+def test_lokr_kohya_fused_qkv_chunked():
+ net = _load_via(F.try_load_lokr, sd_lokr_kohya_qkv())
+ assert net is not None and len(net.modules) == 3
+ for nk, mod in net.modules.items():
+ assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}'
+ assert 0 <= mod.chunk_index < 3 and mod.num_chunks == 3
+ return True
+
+
+def test_lokr_simpletuner_lycoris_format():
+ """Real-world LyCORIS-standalone LoKR with lycoris_ prefix on diffusers paths.
+
+ Critical case: one of the targets is ``add_k_proj`` whose Linear name has
+ embedded underscores. Confirms the loader does NOT do a naive _ -> .
+ expansion on the base path; the underscored form must round-trip through
+ the network_layer_mapping lookup intact.
+ """
+ net = _load_via(F.try_load_lokr, sd_lokr_simpletuner_lycoris_style())
+ assert net is not None and len(net.modules) == 2, f'got {net.modules if net else None}'
+ expected = {
+ 'lora_transformer_transformer_blocks_0_attn_to_q',
+ 'lora_transformer_transformer_blocks_0_attn_add_k_proj',
+ }
+ assert set(net.modules) == expected, f'got {set(net.modules)}'
+ for mod in net.modules.values():
+ assert isinstance(mod, network_lokr.NetworkModuleLokr)
+ return True
+
+
+def test_loha_bfl_non_fused():
+ net = _load_via(F.try_load_loha, sd_loha_bfl_proj())
+ assert net is not None and len(net.modules) == 1
+ nk, mod = next(iter(net.modules.items()))
+ assert isinstance(mod, network_hada.NetworkModuleHada) and not isinstance(mod, network_hada.NetworkModuleHadaChunk), \
+ f'{nk}: type={type(mod).__name__}'
+ return True
+
+
+def test_loha_kohya_fused_qkv_chunked():
+ net = _load_via(F.try_load_loha, sd_loha_kohya_qkv())
+ assert net is not None and len(net.modules) == 3
+ for nk, mod in net.modules.items():
+ assert isinstance(mod, network_hada.NetworkModuleHadaChunk), f'{nk}: type={type(mod).__name__}'
+ assert 0 <= mod.chunk_index < 3 and mod.num_chunks == 3
+ return True
+
+
+def test_oft_kohya_non_fused():
+ net = _load_via(F.try_load_oft, sd_oft_kohya_proj())
+ assert net is not None and len(net.modules) == 1
+ _, mod = next(iter(net.modules.items()))
+ assert isinstance(mod, network_oft.NetworkModuleOFT)
+ assert mod.is_kohya is True and mod.constraint is not None
+ return True
+
+
+def test_oft_lycoris_no_npe():
+ """LyCORIS OFT (oft_diag) loads and calc_updown completes without
+ accessing self.constraint, which is None in the LyCORIS path.
+ """
+ net = _load_via(F.try_load_oft, sd_oft_lycoris_proj())
+ assert net is not None and len(net.modules) == 1
+ _, mod = next(iter(net.modules.items()))
+ assert isinstance(mod, network_oft.NetworkModuleOFT)
+ assert mod.is_kohya is False and mod.constraint is None
+ target = mod.sd_module.weight # real (HIDDEN, QKV_OUT) tensor from the mock
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'LyCORIS OFT updown')
+ assert_finite(updown, 'LyCORIS OFT updown')
+ return True
+
+
+def test_oft_fused_qkv_skipped():
+ """OFT on fused QKV must be dropped (with warning) — no chunk class exists."""
+ net = _load_via(F.try_load_oft, sd_oft_kohya_qkv_skipped())
+ # Loader returns None when no modules bound (skipped + nothing else).
+ assert net is None or len(net.modules) == 0
+ return True
+
+
+def test_boft_loads_and_applies():
+ """BOFT files (4-D oft_blocks) bind to NetworkModuleBOFT via the
+ try_load_oft dispatch. Verifies shape attrs derive correctly from
+ the 4-D tensor and the butterfly-cascade calc_updown returns a
+ finite delta of host weight shape.
+ """
+ net = _load_via(F.try_load_oft, sd_boft_butterfly())
+ assert net is not None and len(net.modules) == 1, \
+ f'BOFT did not load (got {net.modules if net else None})'
+ mod = next(iter(net.modules.values()))
+ assert isinstance(mod, network_boft.NetworkModuleBOFT), \
+ f'expected NetworkModuleBOFT, got {type(mod).__name__}'
+ assert isinstance(mod, network_boft.NetworkModuleBOFT) and not isinstance(mod, network_oft.NetworkModuleOFT), \
+ 'BOFT must not be a subclass of NetworkModuleOFT (separate algorithm)'
+ # BOFT-specific shape attributes resolved correctly
+ assert mod.boft_m == 4 and mod.block_num == 8 and mod.block_size == 16, \
+ f'shape attrs: boft_m={mod.boft_m} block_num={mod.block_num} block_size={mod.block_size}'
+ # End-to-end calc_updown via the butterfly cascade
+ target = mod.sd_module.weight # (HIDDEN=128, QKV_OUT=128) for to_out.0
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'BOFT updown')
+ assert_finite(updown, 'BOFT updown')
+ return True
+
+
+def test_ia3_basic_load():
+ net = _load_via(F.try_load_ia3, sd_ia3_proj())
+ assert net is not None and len(net.modules) == 1
+ _, mod = next(iter(net.modules.items()))
+ assert isinstance(mod, network_ia3.NetworkModuleIa3)
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'IA3 updown')
+ return True
+
+
+def test_ia3_fused_qkv_skipped():
+ net = _load_via(F.try_load_ia3, sd_ia3_qkv_skipped())
+ assert net is None or len(net.modules) == 0
+ return True
+
+
+def test_glora_basic_load_and_alpha():
+ """GLoRA loads, self.dim is populated from w1b.shape[0], and
+ calc_scale returns alpha/dim rather than the 1.0 fallback.
+ """
+ state_dict = sd_glora_proj()
+ net = _load_via(F.try_load_glora, state_dict)
+ assert net is not None and len(net.modules) == 1
+ mod = next(iter(net.modules.values()))
+ assert isinstance(mod, network_glora.NetworkModuleGLora)
+ assert mod.dim is not None, 'GLoRA dim not set (regression)'
+ expected_dim = state_dict['lora_unet_double_blocks_1_img_attn_proj.b1.weight'].shape[0]
+ assert mod.dim == expected_dim, f'GLoRA dim {mod.dim} != {expected_dim}'
+ # alpha=RANK_LORA (8), dim=expected_dim (8) → calc_scale = 1.0; instructive but matches kohya convention.
+ expected_alpha = float(RANK_LORA)
+ assert mod.alpha == expected_alpha, f'alpha={mod.alpha}'
+ expected_scale = expected_alpha / expected_dim
+ actual_scale = mod.calc_scale()
+ assert abs(actual_scale - expected_scale) < 1e-6, f'calc_scale={actual_scale} expected {expected_scale} (alpha/dim)'
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'GLoRA updown')
+ assert_finite(updown, 'GLoRA updown')
+ return True
+
+
+def test_norm_loader_local_stamping():
+ """Regression for the loader-local stamping behavior in try_load_norm.
+
+ lora_convert.assign_network_names_to_compvis_modules deliberately skips
+ setting module.network_layer_name on transformer norm modules. The Norm
+ loader must stamp them itself for any module it actually binds, so the
+ activation loop (which checks network_layer_name is not None) can apply
+ the delta. Verify both: pre-loader the norm module has no
+ network_layer_name; post-loader the bound target does.
+ """
+ sd_model = install_mock_pipe()
+ norm_module = sd_model.pipe.transformer.transformer_blocks[0].attn.norm_q
+ assert not getattr(norm_module, 'network_layer_name', None), \
+ 'mock norm should not have network_layer_name pre-load'
+ # First call assign_network_names manually to mirror what resolve_mapping does.
+ from modules.lora import lora_convert
+ lora_convert.assign_network_names_to_compvis_modules(sd_model.pipe)
+ # The carve-out skipped stamping for norm modules.
+ assert not getattr(norm_module, 'network_layer_name', None), \
+ 'norm module should still be unstamped after assign_network_names (lora_convert.py:502 carve-out)'
+ # Now run the loader.
+ state_dict = sd_norm_peft_attn_norm_q()
+ with TempLora(state_dict, name='norm_test') as nod:
+ net = F.try_load_norm('norm_test', nod, lora_scale=1.0)
+ assert net is not None and len(net.modules) == 1, 'Norm loader produced no modules'
+ mod = next(iter(net.modules.values()))
+ assert isinstance(mod, network_norm.NetworkModuleNorm)
+ # Loader-local stamping: the bound module now has network_layer_name set.
+ assert getattr(norm_module, 'network_layer_name', None) == 'lora_transformer_transformer_blocks_0_attn_norm_q', \
+ f'Norm loader did not stamp network_layer_name (got {getattr(norm_module, "network_layer_name", None)!r})'
+ # calc_updown returns updown + ex_bias for w_norm + b_norm.
+ target = norm_module.weight
+ updown, ex_bias = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'Norm updown')
+ assert ex_bias is not None, 'Norm ex_bias should be present when b_norm is set'
+ assert_shape(ex_bias, tuple(target.shape), 'Norm ex_bias')
+ return True
+
+
+def test_full_basic_load():
+ state_dict = sd_full_proj()
+ net = _load_via(F.try_load_full, state_dict)
+ assert net is not None and len(net.modules) == 1
+ mod = next(iter(net.modules.values()))
+ assert isinstance(mod, network_full.NetworkModuleFull)
+ target = mod.sd_module.weight
+ updown, ex_bias = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape), 'Full updown')
+ assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,), f'Full ex_bias shape {ex_bias.shape if ex_bias is not None else None}'
+ return True
+
+
+# ============================================================
+# Tests — calc_updown sanity per family
+# ============================================================
+
+CAT_MATH = category('math')
+
+
+def _instantiate_first(net):
+ """Convenience: pull the only module out of a single-module Network."""
+ assert net is not None and len(net.modules) == 1
+ return next(iter(net.modules.values()))
+
+
+def test_lora_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_lora, sd_lora_bfl_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_lokr_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_lokr, sd_lokr_bfl_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_lokr_chunk_calc_updown_shape():
+ """LokrChunk produces a (QKV_OUT, HIDDEN) chunk from a (3*QKV_OUT, HIDDEN) full kron."""
+ net = _load_via(F.try_load_lokr, sd_lokr_kohya_qkv())
+ for mod in net.modules.values():
+ target = mod.sd_module.weight # split projection: shape (HIDDEN, QKV_OUT) for to_q etc.
+ updown, _ = mod.calc_updown(target)
+ # `target.shape` is (out=QKV_OUT, in=HIDDEN) for the diffusers Linear; the chunk class
+ # returns the chunk of kron(w1,w2) sliced along dim 0 (the fused-output axis).
+ assert_shape(updown, (QKV_OUT, HIDDEN))
+ assert_finite(updown)
+ return True
+
+
+def test_loha_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_loha, sd_loha_bfl_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_loha_chunk_calc_updown_shape():
+ """HadaChunk produces a (QKV_OUT, HIDDEN) chunk per Q/K/V."""
+ net = _load_via(F.try_load_loha, sd_loha_kohya_qkv())
+ for mod in net.modules.values():
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, (QKV_OUT, HIDDEN))
+ assert_finite(updown)
+ return True
+
+
+def test_oft_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_oft, sd_oft_kohya_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_ia3_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_ia3, sd_ia3_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ # IA3 returns target * w; output shape matches target.
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_glora_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_glora, sd_glora_proj()))
+ target = mod.sd_module.weight
+ updown, _ = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_finite(updown)
+ return True
+
+
+def test_full_calc_updown_shape():
+ mod = _instantiate_first(_load_via(F.try_load_full, sd_full_proj()))
+ target = mod.sd_module.weight
+ updown, ex_bias = mod.calc_updown(target)
+ assert_shape(updown, tuple(target.shape))
+ assert_shape(ex_bias, (HIDDEN,))
+ assert_finite(updown)
+ return True
+
+
+# ============================================================
+# Tests — apply path / regressions in shared infra
+# ============================================================
+
+CAT_APPLY = category('apply')
+
+
+def test_ex_bias_accumulation_two_norms():
+ """Stacking two Norm adapters that both produce ex_bias accumulates
+ correctly: network_calc_weights sums their contributions instead of
+ raising on a tensor truthiness check.
+ """
+ install_mock_pipe()
+ norm_module = shared.sd_model.pipe.transformer.transformer_blocks[0].attn.norm_q
+
+ # Build two Norm Networks targeting the same norm module.
+ state_dict_a = {
+ 'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.full((HEAD_DIM,), 0.01),
+ 'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.full((HEAD_DIM,), 0.02),
+ }
+ state_dict_b = {
+ 'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.full((HEAD_DIM,), 0.03),
+ 'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.full((HEAD_DIM,), 0.04),
+ }
+ nets = []
+ for label, sd in (('a', state_dict_a), ('b', state_dict_b)):
+ with TempLora(sd, name=f'norm_{label}') as nod:
+ n = F.try_load_norm(f'norm_{label}', nod, lora_scale=1.0)
+ n.te_multiplier = 1.0
+ n.unet_multiplier = 1.0
+ nets.append(n)
+
+ # Patch loaded_networks for the duration of this test.
+ saved_loaded = list(l_common.loaded_networks)
+ l_common.loaded_networks.clear()
+ l_common.loaded_networks.extend(nets)
+ try:
+ # network_calc_weights iterates loaded_networks and hits the
+ # `if batch_ex_bias:` line on the second adapter.
+ layer_name = norm_module.network_layer_name
+ # Bind layer_name in case prior test cleared it.
+ if not layer_name:
+ from modules.lora import lora_convert
+ lora_convert.assign_network_names_to_compvis_modules(shared.sd_model.pipe)
+ layer_name = 'lora_transformer_transformer_blocks_0_attn_norm_q'
+ norm_module.network_layer_name = layer_name
+ batch_updown, batch_ex_bias = lora_apply.network_calc_weights(norm_module, layer_name)
+ assert batch_updown is not None, 'batch_updown should be present (two w_norm contributions)'
+ assert batch_ex_bias is not None, 'batch_ex_bias should be present (two b_norm contributions)'
+ assert_shape(batch_ex_bias, (HEAD_DIM,))
+ # Two contributions of 0.02 + 0.04 = 0.06 (pre-multiplier).
+ # multiplier()=te_multiplier=1.0 for both; finalize_updown applies it to ex_bias.
+ expected_b = 0.02 + 0.04
+ assert torch.allclose(batch_ex_bias, torch.full_like(batch_ex_bias, expected_b), atol=1e-5), \
+ f'ex_bias accumulated to {batch_ex_bias[0].item():.6f}, expected {expected_b:.6f}'
+ finally:
+ l_common.loaded_networks.clear()
+ l_common.loaded_networks.extend(saved_loaded)
+ return True
+
+
+# ============================================================
+# Test runner
+# ============================================================
+
+
+def run_tests():
+ t0 = time.time()
+
+ log.warning('=== Parsing primitives ===')
+ for fn in [test_parse_key_all_prefixes, test_resolve_targets_qkv_chunking,
+ test_parse_key_peft_wrapper_unwrap, test_parse_key_lycoris_prefix,
+ test_parse_key_bare_diffusers_and_peft_default,
+ test_marker_disambiguation]:
+ run_test(CAT_PARSE, fn)
+
+ log.warning('=== Loaders ===')
+ for fn in [
+ test_lora_kohya_fused_qkv_chunked,
+ test_lora_bfl_non_fused,
+ test_lora_peft_format,
+ test_lora_bare_bfl_format,
+ test_lora_peft_saved_fal_style,
+ test_lora_peft_saved_dreambooth_style,
+ test_lora_peft_saved_diffusers_style,
+ test_lora_dora_threading,
+ test_lokr_bfl_non_fused,
+ test_lokr_kohya_fused_qkv_chunked,
+ test_lokr_simpletuner_lycoris_format,
+ test_loha_bfl_non_fused,
+ test_loha_kohya_fused_qkv_chunked,
+ test_oft_kohya_non_fused,
+ test_oft_lycoris_no_npe,
+ test_oft_fused_qkv_skipped,
+ test_boft_loads_and_applies,
+ test_ia3_basic_load,
+ test_ia3_fused_qkv_skipped,
+ test_glora_basic_load_and_alpha,
+ test_norm_loader_local_stamping,
+ test_full_basic_load,
+ ]:
+ run_test(CAT_LOADER, fn)
+
+ log.warning('=== calc_updown shape sanity ===')
+ for fn in [
+ test_lora_calc_updown_shape,
+ test_lokr_calc_updown_shape,
+ test_lokr_chunk_calc_updown_shape,
+ test_loha_calc_updown_shape,
+ test_loha_chunk_calc_updown_shape,
+ test_oft_calc_updown_shape,
+ test_ia3_calc_updown_shape,
+ test_glora_calc_updown_shape,
+ test_full_calc_updown_shape,
+ ]:
+ run_test(CAT_MATH, fn)
+
+ log.warning('=== Apply path / shared-infra regressions ===')
+ for fn in [test_ex_bias_accumulation_two_norms]:
+ run_test(CAT_APPLY, fn)
+
+ elapsed = time.time() - t0
+
+ log.warning('=== Results ===')
+ total_passed = 0
+ total_failed = 0
+ for cat, data in results.items():
+ total_passed += data['passed']
+ total_failed += data['failed']
+ status = 'PASS' if data['failed'] == 0 else 'FAIL'
+ log.info(f' {cat}: {data["passed"]} passed, {data["failed"]} failed [{status}]')
+ log.warning(f'Total: {total_passed} passed, {total_failed} failed in {elapsed:.2f}s')
+ if total_failed:
+ sys.exit(1)
+
+
+if __name__ == '__main__':
+ run_tests()
diff --git a/test/test-generation-api.py b/test/test-generation-api.py
index 1359d59e5..98ec18982 100644
--- a/test/test-generation-api.py
+++ b/test/test-generation-api.py
@@ -417,7 +417,7 @@ class GenerationAPITest:
self.skip(f'param_{name}', 'baseline generation failed')
return
- data, elapsed = self._txt2img(params)
+ data, _elapsed = self._txt2img(params)
if 'error' in data:
self.record(False, f'param_{name}', f"generation error: {data}")
return
@@ -505,7 +505,7 @@ class GenerationAPITest:
# Vignette: corners should be darker than baseline corners
def check_vignette(base, result, _data):
- h, w = base.shape[:2]
+ h, _w = base.shape[:2]
corner_size = h // 8
base_corners = np.concatenate([
base[:corner_size, :corner_size].flatten(),
diff --git a/wiki b/wiki
index 070218735..099acbacc 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit 07021873564383be3456b673ad27693db31a8d9a
+Subproject commit 099acbacc662634b22732aebcf4f04691ae2f26a