fix issues and further cleanup

Author: damian0815

ldm/invoke/conditioning.py

@@ -15,6 +15,7 @@
 from compel.prompt_parser import FlattenedPrompt, Blend, Fragment, CrossAttentionControlSubstitute, PromptParser, \
     Conjunction
 from .devices import torch_dtype
+from .generator.diffusers_pipeline import StableDiffusionGeneratorPipeline
 from ..models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
 from ldm.invoke.globals import Globals
 
@@ -32,7 +33,7 @@ def __call__(self, *args, **kwargs):
 
 
 def get_uc_and_c_and_ec(prompt_string,
-                        model: StableDiffusionPipeline,
+                        model: StableDiffusionGeneratorPipeline,
                         log_tokens=False, skip_normalize_legacy_blend=False):
     # lazy-load any deferred textual inversions.
     # this might take a couple of seconds the first time a textual inversion is used.
@@ -75,7 +76,7 @@ def get_uc_and_c_and_ec(prompt_string,
     # some LoRA models also mess with the text encoder, so they must be active while compel builds conditioning tensors
     lora_conditioning_ec = InvokeAIDiffuserComponent.ExtraConditioningInfo(tokens_count_including_eos_bos=tokens_count,
                                                                                             lora_conditions=lora_conditions)
-    with InvokeAIDiffuserComponent.custom_attention_context(model,
+    with InvokeAIDiffuserComponent.custom_attention_context(model.unet,
                                                             extra_conditioning_info=lora_conditioning_ec,
                                                             step_count=-1):
         c, options = compel.build_conditioning_tensor_for_prompt_object(positive_prompt)

ldm/models/diffusion/cross_attention_control.py

@@ -288,7 +288,7 @@ def get_invokeai_attention_mem_efficient(self, q, k, v):
         return self.einsum_op_tensor_mem(q, k, v, 32)
 
 
-def setup_cross_attention_control_attention_processors(model, context: Context):
+def setup_cross_attention_control_attention_processors(unet: UNet2DConditionModel, context: Context):
     """
     Inject attention parameters and functions into the passed in model to enable cross attention editing.
 
@@ -314,7 +314,6 @@ def setup_cross_attention_control_attention_processors(model, context: Context):
 
     context.cross_attention_mask = mask.to(device)
     context.cross_attention_index_map = indices.to(device)
-    unet = model
     old_attn_processors = unet.attn_processors
     if torch.backends.mps.is_available():
         # see note in StableDiffusionGeneratorPipeline.__init__ about borked slicing on MPS

ldm/models/diffusion/shared_invokeai_diffusion.py

@@ -5,7 +5,8 @@
 
 import numpy as np
 import torch
-from diffusers.models.cross_attention import CrossAttnProcessor
+
+from diffusers import UNet2DConditionModel
 from typing_extensions import TypeAlias
 
 from ldm.invoke.globals import Globals
@@ -86,33 +87,36 @@ def __init__(
     @classmethod
     @contextmanager
     def custom_attention_context(
-        clss, model, extra_conditioning_info: Optional[ExtraConditioningInfo], step_count: int
+        clss,
+        unet: UNet2DConditionModel, # note: also may futz with the text encoder depending on requested LoRAs
+        extra_conditioning_info: Optional[ExtraConditioningInfo],
+        step_count: int
     ):
         old_attn_processors = None
         if extra_conditioning_info and (
             extra_conditioning_info.wants_cross_attention_control
             | extra_conditioning_info.has_lora_conditions
         ):
-            old_attn_processors = model.attn_processors
+            old_attn_processors = unet.attn_processors
             # Load lora conditions into the model
             if extra_conditioning_info.has_lora_conditions:
                 for condition in extra_conditioning_info.lora_conditions:
-                    condition(model)
+                    condition() # target model is stored in condition state for some reason
             if extra_conditioning_info.wants_cross_attention_control:
                 cross_attention_control_context = Context(
                     arguments=extra_conditioning_info.cross_attention_control_args,
                     step_count=step_count,
                 )
                 setup_cross_attention_control_attention_processors(
-                    model,
+                    unet,
                     cross_attention_control_context,
                 )
 
         try:
             yield None
         finally:
             if old_attn_processors is not None:
-                model.set_attn_processor(old_attn_processors)
+                unet.set_attn_processor(old_attn_processors)
             if extra_conditioning_info and extra_conditioning_info.has_lora_conditions:
                 for lora_condition in extra_conditioning_info.lora_conditions:
                     lora_condition.unload()

ldm/modules/kohya_lora_manager.py

@@ -86,7 +86,7 @@ def forward(self, lora, input_h, output):
             rebuild1 = torch.einsum('i j k l, j r, i p -> p r k l', self.t1, self.w1_b, self.w1_a)
             rebuild2 = torch.einsum('i j k l, j r, i p -> p r k l', self.t2, self.w2_b, self.w2_a)
             weight = rebuild1 * rebuild2
-        
+
         bias = self.bias if self.bias is not None else 0
         return output + op(
             *input_h,

ldm/modules/lora_manager.py

@@ -1,5 +1,7 @@
 import os
 from pathlib import Path
+
+from diffusers import UNet2DConditionModel, StableDiffusionPipeline
 from ldm.invoke.globals import global_lora_models_dir
 from .kohya_lora_manager import KohyaLoraManager
 from typing import Optional, Dict
@@ -8,20 +10,29 @@ class LoraCondition:
     name: str
     weight: float
 
-    def __init__(self, name, weight: float = 1.0, kohya_manager: Optional[KohyaLoraManager]=None):
+    def __init__(self,
+                 name,
+                 weight: float = 1.0,
+                 unet: UNet2DConditionModel=None,  # for diffusers format LoRAs
+                 kohya_manager: Optional[KohyaLoraManager]=None,  # for KohyaLoraManager-compatible LoRAs
+                 ):
         self.name = name
         self.weight = weight
         self.kohya_manager = kohya_manager
+        self.unet = unet
 
-    def __call__(self, model):
+    def __call__(self):
         # TODO: make model able to load from huggingface, rather then just local files
         path = Path(global_lora_models_dir(), self.name)
         if path.is_dir():
-            if model.load_attn_procs:
+            if not self.unet:
+                print(f"   ** Unable to load diffusers-format LoRA {self.name}: unet is None")
+                return
+            if self.unet.load_attn_procs:
                 file = Path(path, "pytorch_lora_weights.bin")
                 if file.is_file():
                     print(f">> Loading LoRA: {path}")
-                    model.load_attn_procs(path.absolute().as_posix())
+                    self.unet.load_attn_procs(path.absolute().as_posix())
                 else:
                     print(f"   ** Unable to find valid LoRA at: {path}")
             else:
@@ -37,15 +48,16 @@ def unload(self):
             self.kohya_manager.unload_applied_lora(self.name)
 
 class LoraManager:
-    def __init__(self, pipe):
+    def __init__(self, pipe: StableDiffusionPipeline):
         # Kohya class handles lora not generated through diffusers
         self.kohya = KohyaLoraManager(pipe, global_lora_models_dir())
+        self.unet = pipe.unet
 
     def set_loras_conditions(self, lora_weights: list):
         conditions = []
         if len(lora_weights) > 0:
             for lora in lora_weights:
-                conditions.append(LoraCondition(lora.model, lora.weight, self.kohya))
+                conditions.append(LoraCondition(lora.model, lora.weight, self.unet, self.kohya))
 
         if len(conditions) > 0:
             return conditions
@@ -63,4 +75,4 @@ def list_loras(self)->Dict[str, Path]:
                 if suffix in [".ckpt", ".pt", ".safetensors"]:
                     models_found[name]=Path(root,x)
         return models_found
-            
+