Fixed problems with turning off FreeU in some scenarios

Author: MoonRide303

fooocus_version.py

@@ -1,2 +1,2 @@
-version = '2.0.78.3 MRE'
+version = '2.0.78.4 MRE'
 full_version = 'Fooocus ' + version

modules/core.py

@@ -187,7 +187,8 @@ def detect_edge(image, low_threshold, high_threshold):
 @torch.no_grad()
 @torch.inference_mode()
 def freeu(model, b1, b2, s1, s2):
-    return opFreeU.patch(model=model, b1=b1, b2=b2, s1=s1, s2=s2)[0]
+    unet = opFreeU.patch(model=model.unet, b1=b1, b2=b2, s1=s1, s2=s2)[0]
+    return StableDiffusionModel(unet=unet, clip=model.clip, vae=model.vae, clip_vision=model.clip_vision)
 
 
 @torch.no_grad()

modules/default_pipeline.py

@@ -130,9 +130,10 @@ def patch_base(loras, freeu, b1, b2, s1, s2):
         assert os.path.exists(filename), 'Lora file not found!'
 
         model = core.load_sd_lora(model, filename, strength_model=weight, strength_clip=weight)
-    xl_base_patched = model
     if freeu:
-        xl_base_patched.unet = core.freeu(xl_base_patched.unet, b1, b2, s1, s2)
+        xl_base_patched = core.freeu(model, b1, b2, s1, s2)
+    else:
+        xl_base_patched = model
     xl_base_patched_hash = str(loras + [freeu, b1, b2, s1, s2])
     print(f'LoRAs loaded: {loras}')
     if freeu:

update_log_mre.md

@@ -1,3 +1,7 @@
+### 2.0.78.4 MRE
+
+* Fixed problems with turning off FreeU in some scenarios.
+
 ### 2.0.78.3 MRE
 
 * Added limited support for non-SDXL models (no refiner, Control-LoRAs, Revision, inpainting, outpainting).