Replace CLIP monkey-patching with add_object_patch

README.md

@@ -15,9 +15,6 @@ Like any other model patch:
 ![combined_image_new](https://github.com/Extraltodeus/Stable-Diffusion-temperature-settings/assets/15731540/12034834-43d0-44a5-a603-6c87d1bc6e5d)
 
 
-Note: The CLIP temperature patches the attention within the VRAM/RAM on the connected clip model. The modification ignores connections but is only linked to what is loaded. To get default behavior set the temperature at one or reload the model.
-
-
 # Patreon
 
 Provide an incentive to contributors:

nodes.py

@@ -1,7 +1,10 @@
+from functools import partial
 import torch
 import math
 import types
 import comfy.model_management
+from comfy.model_patcher import ModelPatcher
+from comfy.sd import CLIP
 
 EPSILON = 1e-16
 
@@ -36,7 +39,7 @@ class temperature_patcher():
         )
 
         scale = 1 / (math.sqrt(q.size(-1)) * self.temperature)
-        
+
         out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale)
         if should_scale(self.model_name, self.layer_name,q.size(-2)):
             ldim = SD_layer_dims[self.model_name][self.layer_name]
@@ -69,7 +72,7 @@ class UnetTemperaturePatch:
 
     CATEGORY = "model_patches/Temperature"
 
-    def patch(self, model, Temperature, Attention, Dynamic_Scale_Attention, eval_string="", **kwargs):
+    def patch(self, model: ModelPatcher, Temperature, Attention, Dynamic_Scale_Attention, eval_string="", **kwargs):
         if Dynamic_Scale_Attention and str(model.size) in models_by_size:
             model_name = models_by_size[str(model.size)]
             print(f"Model detected for scaling: {model_name}")
@@ -100,15 +103,17 @@ class CLIPTemperaturePatch:
         return {"required": { "clip": ("CLIP",),
                               "Temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                               }}
-    
+
     RETURN_TYPES = ("CLIP",)
     FUNCTION = "patch"
     CATEGORY = "model_patches/Temperature"
-    
-    def patch(self, clip, Temperature):
+
+    def patch(self, clip: CLIP, Temperature):
+        c = clip.clone()
+
         def custom_optimized_attention(device, mask=None, small_input=True):
             return temperature_patcher(Temperature).pytorch_attention_with_temperature
-        
+
         def new_forward(self, x, mask=None, intermediate_output=None):
             optimized_attention = custom_optimized_attention(x.device, mask=mask is not None, small_input=True)
 
@@ -123,11 +128,10 @@ class CLIPTemperaturePatch:
                     intermediate = x.clone()
             return x, intermediate
 
-        clip_encoder_instance = clip.cond_stage_model.clip_l.transformer.text_model.encoder
-        clip_encoder_instance.forward = types.MethodType(new_forward, clip_encoder_instance)
+        if getattr(c.patcher.model, "clip_g", None) is not None:
+            c.patcher.add_object_patch("clip_g.transformer.text_model.encoder.forward", partial(new_forward, c.patcher.model.clip_g.transformer.text_model.encoder))
 
-        if getattr(clip.cond_stage_model, f"clip_g", None) is not None:
-            clip_encoder_instance_g = clip.cond_stage_model.clip_g.transformer.text_model.encoder
-            clip_encoder_instance_g.forward = types.MethodType(new_forward, clip_encoder_instance_g)
-        
-        return (clip,)
+        if getattr(c.patcher.model, "clip_l", None) is not None:
+            c.patcher.add_object_patch("clip_l.transformer.text_model.encoder.forward", partial(new_forward, c.patcher.model.clip_l.transformer.text_model.encoder))
+
+        return (c,)