diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 8df1f98..7aa90d9 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -29,7 +29,8 @@ from transformers import ( ) from optimum.quanto import freeze, QTensor -from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig +from toolkit.lora_special import LoRASpecialNetwork from toolkit.models.base_model import BaseModel from toolkit.basic import flush from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds @@ -115,7 +116,9 @@ def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict: # Treat as a hub repo id. When no filename is given, derive it from the repo # name's trailing segment (e.g. "krea/Krea-2-Raw" -> "raw.safetensors", # "krea/Krea-2-Turbo" -> "turbo.safetensors"). - fname = filename or (name_or_path.split("/")[-1].split("-")[-1].lower() + ".safetensors") + fname = filename or ( + name_or_path.split("/")[-1].split("-")[-1].lower() + ".safetensors" + ) try: path = huggingface_hub.hf_hub_download( repo_id=name_or_path, filename=fname, token=HF_TOKEN @@ -230,12 +233,101 @@ class Krea2Model(BaseModel): vae.requires_grad_(False) return vae + def load_training_adapter(self, transformer: SingleStreamDiT): + self.print_and_status_update("Loading assistant LoRA") + lora_path = self.model_config.assistant_lora_path + if not os.path.exists(lora_path): + # assume it is a hub path + lora_splits = lora_path.split("/") + if len(lora_splits) != 3: + raise ValueError( + f"Assistant LoRA path {lora_path} is not a valid local path or hub path." + ) + repo_id = "/".join(lora_splits[:2]) + filename = lora_splits[2] + try: + lora_path = huggingface_hub.hf_hub_download( + repo_id=repo_id, + filename=filename, + token=HF_TOKEN, + ) + # upgrade path to the local download + self.model_config.assistant_lora_path = lora_path + except Exception as e: + raise ValueError( + f"Failed to download assistant LoRA from {lora_path}: {e}" + ) + # load the adapter and merge it in. We will inference with a -1.0 multiplier so the adapter effects only work during training. + lora_state_dict = load_file(lora_path) + # detect the LoRA rank from the first down-projection weight. + dim_key = next(k for k in lora_state_dict if k.endswith("lora_A.weight")) + dim = int(lora_state_dict[dim_key].shape[0]) + + new_sd = {} + for key, value in lora_state_dict.items(): + new_key = key.replace("diffusion_model.", "transformer.") + new_sd[new_key] = value + lora_state_dict = new_sd + + network_config = { + "type": "lora", + "linear": dim, + "linear_alpha": dim, + "transformer_only": True, + } + + network_config = NetworkConfig(**network_config) + LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer" + network = LoRASpecialNetwork( + text_encoder=None, + unet=transformer, + lora_dim=network_config.linear, + multiplier=1.0, + alpha=network_config.linear_alpha, + train_unet=True, + train_text_encoder=False, + network_config=network_config, + network_type=network_config.type, + transformer_only=network_config.transformer_only, + is_transformer=True, + target_lin_modules=self.target_lora_modules, + is_assistant_adapter=True, + is_ara=True, + ) + network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True) + self.print_and_status_update("Merging in assistant LoRA") + network.force_to(self.device_torch, dtype=self.torch_dtype) + network._update_torch_multiplier() + network.load_weights(lora_state_dict) + + network.merge_in(merge_weight=1.0) + + # mark it as not merged so inference ignores it. + network.is_merged_in = False + + # add the assistant so sampler will activate it while sampling + self.assistant_lora: LoRASpecialNetwork = network + + # deactivate lora during training + self.assistant_lora.multiplier = -1.0 + self.assistant_lora.is_active = False + + # tell the model to invert assistant on inference since we want remove lora effects + self.invert_assistant_lora = True + def load_model(self): dtype = self.torch_dtype self.print_and_status_update("Loading Krea 2 model") transformer = self._load_transformer() + # load assistant lora if specified + if self.model_config.assistant_lora_path is not None: + self.load_training_adapter(transformer) + # set qtype to be float8 if it is qfloat8 + if self.model_config.qtype == "qfloat8": + self.model_config.qtype = "float8" + if self.model_config.quantize: self.print_and_status_update("Quantizing transformer") quantize_model(self, transformer) diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 0979c96..74c7335 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -1041,7 +1041,7 @@ export const modelArchs: ModelArch[] = [ }, { name: 'krea2', - label: 'Krea 2 (K2)', + label: 'Krea 2 (raw)', group: 'image', defaults: { 'config.process[0].model.name_or_path': ['krea/Krea-2-Raw', defaultNameOrPath], @@ -1060,6 +1060,34 @@ export const modelArchs: ModelArch[] = [ 'model.layer_offloading', ], }, + { + name: 'krea2:turbo', + label: 'Krea 2 Turbo (w/ Training Adapter)', + group: 'image', + defaults: { + 'config.process[0].model.name_or_path': ['krea/Krea-2-Turbo', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + 'config.process[0].model.assistant_lora_path': [ + 'ostris/krea2_turbo_training_adapter/krea2_turbo_training_adapter_v1.safetensors', + undefined, + ], + 'config.process[0].sample.guidance_scale': [1, 4], + 'config.process[0].sample.sample_steps': [8, 25], + }, + disableSections: [ + 'network.conv', + ], + additionalSections: [ + 'model.low_vram', + 'model.layer_offloading', + 'model.assistant_lora_path' + ], + }, { name: 'boogu_image', label: 'Boogu Image', diff --git a/version.py b/version.py index 567fcdc..bf58814 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.17" +VERSION = "0.10.18"