New Model: Krea 2 - #1568
Merged
Merged
Conversation
Replace the per-model key_set machinery with a single convert()-native pipeline shared across all model types. Each model owns its diffusers<->original mapping and declares its text encoders and fusion groups once; savers and loaders derive every format from those. Save formats reworked into a small enum (KOHYA, LEGACY, ORIGINAL, COMFY, DIFFUSERS) with real kohya-ss native+fused output, a frozen LEGACY that absorbs the old format, and full-model COMFY_TRANSFORMER support. Loading auto-detects the source format and normalizes to canonical keys with per-model legacy hooks. Migrates SD1.5/2.x, SDXL, SD3/3.5, Flux, Flux2, Chroma, PixArt, Sana, Qwen, Ernie, Z-Image, HiDream, HunyuanVideo, Stable Cascade and Wuerstchen, plus the standalone LoRA/checkpoint conversion tool. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Resolves the Ctk/PySide6 view-controller split against the LoRA/full-model output-format rework: LoRAModule.py keeps the fusion check while dropping the rank check (upstream Nerogar#1549); ModelTab.py's split becomes BaseModelTabView.py + ModelTabController.py, with the output-format selection logic living in the controller (get_output_formats), matching the TopBarController pattern.
…nto Nerogar:merge Adds AttentionMechanism (SDP/FLASH) selection, restoring the explicit attention-backend choice removed when xformers was dropped. Rebased from PR Nerogar#1227 (originally against master) onto Nerogar:merge; the UI addition was re-targeted at BaseTrainingTabView (the ctk/Qt6 split from PR Nerogar#1566) since the old TrainingTab.py no longer exists on this base. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…verage - Add AttentionMechanism.CUDNN (diffusers "_native_cudnn" backend). - Make _set_attention_backend's mask argument required (no default) so every call site has to make an explicit decision; drop the unused varlen parameter (FLASH hard-raises on any mask, it doesn't "maybe fail"). - Roll out to the models the original PR missed: Ernie (mask=True, its transformer builds a real attention mask internally from text_lens), Flux2 (mask=False), Wuerstchen (mask=False, StableCascadeUNet.forward has no attention_mask parameter at all). - Switch the UI dropdown to options_kv (same pattern as LoraTabController.get_peft_types()) with human-readable labels sourced from TrainingTabController.get_attention_mechanisms(), keeping the stored/serialized enum identity unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Removed an extra newline before setting the attention backend.
Adds Krea 2 (transformer + Qwen3-VL text encoder + Qwen-Image VAE) as a new model type: sampling pipeline, LoRA and full finetune training, checkpoint loading/saving (diffusers + native raw.safetensors formats), UI wiring, and a starter 16GB LoRA preset. Text-encoder training is not supported (matches Z-Image/Ernie). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This comment was marked as resolved.
This comment was marked as resolved.
FusedModuleGroup._leaf_forward_N called the inner peft module's forward once per leaf, and that forward's orig_forward is the synthetic fused base recomputing all N leaves' base output every time -- N real base computes per leaf, N^2 total for a group of N. Add PeftBase.delta_forward, an optional hook returning just the adapter's own contribution (the term added to orig_forward(x)) without touching the base. LoRAModule and LoHaModule implement it (their deltas never read the base weight); FusedModuleGroup uses it, when available, to add each leaf's real, unfused base once instead. DoRA/OFT/LoKr keep the slower generic path (delta_forward defaults to None) since their forwards recompose the base weight itself and aren't expressible this way. Also replace the functools.partial(self._leaf_forward, leaf_index) hook with four fixed _leaf_forward_0..3 methods: torch.compile guards on leaf.forward's function identity, and a partial rebuilt on every hook_to_module() call was a fresh object each time, forcing a recompile on every hook/unhook cycle. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as off-topic.
This comment was marked as off-topic.
Remove non-explanatory/out-of-place comments, drop the untrue Cascade bracket note on LEGACY_LORA, and revert unnecessary learning_rate notation-only changes in a few presets. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…uess Mirrors the saver's _convert_legacy opt-in default: a model gets LEGACY load support only by explicitly declaring its historical layout, not by inheriting a generic reconstruction that happens to be right for most models but silently wrong for any that isn't. Flux/Flux2/HunyuanVideo/SD3 opt in via the extracted _mixture_legacy_conversion() helper; Sana's now- redundant None override is removed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Trims explanatory-only comments, converts ModelType's LEGACY/ORIGINAL_TRANSFORMER format support from denylists to allow-lists (new models default to unsupported until a saver implements them), drops the qkv-group matching's dependency on a literal "i" placeholder name (matched_leaf_groups works with any/no placeholder), removes dead fuse_qkv/fuse_split duplication in favor of a single variadic fuse(), and removes the _denoising_body_conversion identity wrapper.
…ression - Add required base-model-name field to the convert tool for LoRA/embedding conversion - Fix convert_model.py CLI crashing on any LoRA/embedding conversion - Inline the unused _check_fusion_match indirection in LoRASaverMixin - Fix Flux2 LEGACY LoRA output regression and remove leftover debug print
Squashed review fixes: centralize the LoRA param allowlist and derive SUPPORTED_PARAM_PREFIXES from FACTOR_PREFIXES, filter absent text encoders, drop dead guards, fix the DoRA chunk-swap, and trim duplicate comments. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
# Conflicts: # modules/module/FusedModule.py
Remove dead text-encoder-LoRA/train-dtype fields and now-pointless text_encoder-is-None guards (TE is always loaded and never trained in Krea2), drop redundant comments, simplify a few one-liners, and add TODO markers for code that should eventually be shared with other models. Also tighten supported_lora_formats/supported_full_model_formats to use is_flux_1()/is_flux_2() explicitly instead of the version- spanning is_flux(), so a future Flux variant doesn't silently inherit unverified legacy-format support.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
contains #1563
Test plan
pre-commit run --all-filespassesAI assistance