Update forge_freeu.py

extensions-builtin/sd_forge_freeu/scripts/forge_freeu.py

@@ -16,7 +16,7 @@ def Fourier_filter(x, threshold, scale):
     return x_filtered.to(x.dtype)
 
 
-def patch(model, b1, b2, s1, s2):
+def freeu_v2(model, b1, b2, s1, s2):
     model_channels = model.model.model_config.unet_config["model_channels"]
     scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
 
@@ -62,6 +62,12 @@ class FreeUForForge(scripts.Script):
         if not freeu_enabled:
             return
 
+        unet = p.sd_model.forge_objects.unet
+
+        unet = freeu_v2(unet, freeu_b1, freeu_b2, freeu_s1, freeu_s2)
+
+        p.sd_model.forge_objects.unet = unet
+
         # Below codes will add some logs to the texts below the image outputs on UI.
         # The extra_generation_params does not influence results.
         p.extra_generation_params.update(dict(