wip: update ldm_patched

currently issues with calculate_sigmas call

modules/patch_clip.py

@@ -1,5 +1,6 @@
 # Consistent with Kohya/A1111 to reduce differences between model training and inference.
 
+import json
 import os
 import torch
 import ldm_patched.controlnet.cldm
@@ -62,48 +63,42 @@ def patched_encode_token_weights(self, token_weight_pairs):
     return torch.cat(output, dim=-2).to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
 
 
-def patched_SDClipModel__init__(self, max_length=77, freeze=True, layer="last", layer_idx=None,
-                                textmodel_json_config=None, dtype=None, special_tokens=None,
-                                layer_norm_hidden_state=True, **kwargs):
+def patched_SDClipModel__init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
+                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=ldm_patched.modules.clip_model.CLIPTextModel,
+                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, return_projected_pooled=True):  # clip-vit-base-patch32
     torch.nn.Module.__init__(self)
     assert layer in self.LAYERS
 
-    if special_tokens is None:
-        special_tokens = {"start": 49406, "end": 49407, "pad": 49407}
-
     if textmodel_json_config is None:
-        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)),
-                                             "sd1_clip_config.json")
+        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)), "sd1_clip_config.json")
 
-    config = CLIPTextConfig.from_json_file(textmodel_json_config)
-    self.num_layers = config.num_hidden_layers
-
-    with use_patched_ops(ops.manual_cast):
-        with modeling_utils.no_init_weights():
-            self.transformer = CLIPTextModel(config)
-
-    if dtype is not None:
-        self.transformer.to(dtype)
+    with open(textmodel_json_config) as f:
+        config = json.load(f)
 
+    self.transformer = model_class(config, dtype, device, ldm_patched.modules.ops.manual_cast)
+    self.num_layers = self.transformer.num_layers
     self.transformer.text_model.embeddings.to(torch.float32)
 
+    self.max_length = max_length
     if freeze:
         self.freeze()
-
-    self.max_length = max_length
     self.layer = layer
     self.layer_idx = None
     self.special_tokens = special_tokens
+
+    # TODO check if necessary
     self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
     self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
-    self.enable_attention_masks = False
+    self.enable_attention_masks = enable_attention_masks
 
     self.layer_norm_hidden_state = layer_norm_hidden_state
+    self.return_projected_pooled = return_projected_pooled
+
     if layer == "hidden":
         assert layer_idx is not None
         assert abs(layer_idx) < self.num_layers
-        self.clip_layer(layer_idx)
-    self.layer_default = (self.layer, self.layer_idx)
+        self.set_clip_options({"layer": layer_idx})
+    self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
 
 
 def patched_SDClipModel_forward(self, tokens):
@@ -122,8 +117,7 @@ def patched_SDClipModel_forward(self, tokens):
                 if tokens[x, y] == max_token:
                     break
 
-    outputs = self.transformer(input_ids=tokens, attention_mask=attention_mask,
-                               output_hidden_states=self.layer == "hidden")
+    outputs = self.transformer(tokens, attention_mask, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
     self.transformer.set_input_embeddings(backup_embeds)
 
     if self.layer == "last":
@@ -131,17 +125,14 @@ def patched_SDClipModel_forward(self, tokens):
     elif self.layer == "pooled":
         z = outputs.pooler_output[:, None, :]
     else:
-        z = outputs.hidden_states[self.layer_idx]
-        if self.layer_norm_hidden_state:
-            z = self.transformer.text_model.final_layer_norm(z)
+        z = outputs[1]
 
-    if hasattr(outputs, "pooler_output"):
-        pooled_output = outputs.pooler_output.float()
-    else:
-        pooled_output = None
-
-    if self.text_projection is not None and pooled_output is not None:
-        pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
+    pooled_output = None
+    if len(outputs) >= 3:
+        if not self.return_projected_pooled and len(outputs) >= 4 and outputs[3] is not None:
+            pooled_output = outputs[3].float()
+        elif outputs[2] is not None:
+            pooled_output = outputs[2].float()
 
     return z.float(), pooled_output