Added support for FreeU

Author: MoonRide303

fooocus_version.py

@@ -1,2 +1,2 @@
-version = '2.0.78.1 MRE'
+version = '2.0.78.2 MRE'
 full_version = 'Fooocus ' + version

launch.py

@@ -24,7 +24,7 @@ def prepare_environment():
     xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.21')
 
     comfy_repo = os.environ.get('COMFY_REPO', "https://github.com/comfyanonymous/ComfyUI")
-    comfy_commit_hash = os.environ.get('COMFY_COMMIT_HASH', "492db2de8db7e082addf131b40adb4a1b7535821")
+    comfy_commit_hash = os.environ.get('COMFY_COMMIT_HASH', "2381d36e6db8e8150e42ff2ede628db5b00ae26f")
 
     print(f"Python {sys.version}")
     print(f"Fooocus version: {fooocus_version.version}")

modules/async_worker.py

@@ -77,6 +77,7 @@ def handler(task):
         revision_strength_3, revision_strength_4, same_seed_for_all, output_format, \
         control_lora_canny, canny_edge_low, canny_edge_high, canny_start, canny_stop, canny_strength, canny_model, \
         control_lora_depth, depth_start, depth_stop, depth_strength, depth_model, use_expansion, \
+        freeu, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
         input_image_checkbox, current_tab, \
         uov_method, uov_input_image, outpaint_selections, inpaint_input_image, \
         input_gallery, revision_gallery, keep_input_names = task
@@ -302,7 +303,12 @@ def handler(task):
         pipeline.refresh_everything(
             refiner_model_name=refiner_model_name,
             base_model_name=base_model_name,
-            loras=loras)
+            loras=loras,
+            freeu=freeu,
+            b1=freeu_b1,
+            b2=freeu_b2,
+            s1=freeu_s1,
+            s2=freeu_s2)
 
         pipeline.set_clip_skips(base_clip_skip, refiner_clip_skip)
         if revision_mode:
@@ -503,11 +509,16 @@ def callback(step, x0, x, total_steps, y):
                     'base_clip_skip': base_clip_skip, 'refiner_clip_skip': refiner_clip_skip,
                     'base_model': base_model_name, 'refiner_model': refiner_model_name,
                     'l1': l1, 'w1': w1, 'l2': l2, 'w2': w2, 'l3': l3, 'w3': w3,
-                    'l4': l4, 'w4': w4, 'l5': l5, 'w5': w5, 'img2img': img2img_mode, 'revision': revision_mode,
+                    'l4': l4, 'w4': w4, 'l5': l5, 'w5': w5, 'freeu': freeu,
+                    'img2img': img2img_mode, 'revision': revision_mode,
                     'positive_prompt_strength': positive_prompt_strength, 'negative_prompt_strength': negative_prompt_strength,
                     'control_lora_canny': control_lora_canny, 'control_lora_depth': control_lora_depth,
                     'prompt_expansion': use_expansion
                 }
+                if freeu:
+                    metadata |= {
+                        'freeu_b1': freeu_b1, 'freeu_b2': freeu_b2, 'freeu_s1': freeu_s1, 'freeu_s2': freeu_s2
+                    }
                 if img2img_mode:
                     metadata |= {
                         'start_step': start_step, 'denoise': denoise, 'scale': img2img_scale, 'input_image': input_image_filename
@@ -548,6 +559,7 @@ def callback(step, x0, x, total_steps, y):
                         ('CFG & CLIP Skips', (cfg, base_clip_skip, refiner_clip_skip)),
                         ('Base Model', base_model_name),
                         ('Refiner Model', refiner_model_name),
+                        ('FreeU', (freeu, freeu_b1, freeu_b2, freeu_s1, freeu_s2) if freeu else (freeu)),
                         ('Image-2-Image', (img2img_mode, start_step, denoise, img2img_scale, input_image_filename) if img2img_mode else (img2img_mode)),
                         ('Revision', (revision_mode, revision_strength_1, revision_strength_2, revision_strength_3,
                             revision_strength_4, revision_images_filenames) if revision_mode else (revision_mode)),

modules/core.py

@@ -18,6 +18,7 @@
 from comfy.sample import prepare_mask, broadcast_cond, get_additional_models, cleanup_additional_models
 from comfy_extras.nodes_post_processing import ImageScaleToTotalPixels
 from comfy_extras.nodes_canny import Canny
+from comfy_extras.nodes_freelunch import FreeU
 from comfy.model_base import SDXLRefiner
 from comfy.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
 from modules.samplers_advanced import KSamplerBasic, KSamplerWithRefiner
@@ -36,6 +37,7 @@
 opCLIPVisionEncode = CLIPVisionEncode()
 opUnCLIPConditioning = unCLIPConditioning()
 opCanny = Canny()
+opFreeU = FreeU()
 opControlNetApplyAdvanced = ControlNetApplyAdvanced()
 
 
@@ -182,6 +184,12 @@ def detect_edge(image, low_threshold, high_threshold):
     return opCanny.detect_edge(image=image, low_threshold=low_threshold, high_threshold=high_threshold)[0]
 
 
+@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]
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent):

modules/default_pipeline.py

@@ -104,9 +104,9 @@ def refresh_refiner_model(name):
 
 @torch.no_grad()
 @torch.inference_mode()
-def refresh_loras(loras):
+def patch_base(loras, freeu, b1, b2, s1, s2):
     global xl_base, xl_base_patched, xl_base_patched_hash
-    if xl_base_patched_hash == str(loras):
+    if xl_base_patched_hash == str(loras + [freeu, b1, b2, s1, s2]):
         return
 
     model = xl_base
@@ -123,13 +123,16 @@ def refresh_loras(loras):
 
         model = core.load_sd_lora(model, filename, strength_model=weight, strength_clip=weight)
     xl_base_patched = model
-    xl_base_patched_hash = str(loras)
-    print(f'LoRAs loaded: {xl_base_patched_hash}')
+    if freeu:
+        xl_base_patched.unet = core.freeu(xl_base_patched.unet, b1, b2, s1, s2)
+    xl_base_patched_hash = str(loras + [freeu, b1, b2, s1, s2])
+    print(f'LoRAs loaded: {loras}')
+    if freeu:
+        print(f'FreeU applied: {[b1, b2, s1, s2]}')
 
     return
 
 
-
 @torch.no_grad()
 @torch.inference_mode()
 def refresh_clip_vision():
@@ -281,7 +284,7 @@ def clear_all_caches():
 
 @torch.no_grad()
 @torch.inference_mode()
-def refresh_everything(refiner_model_name, base_model_name, loras):
+def refresh_everything(refiner_model_name, base_model_name, loras, freeu, b1, b2, s1, s2):
     refresh_refiner_model(refiner_model_name)
     if xl_refiner is not None:
         virtual_memory.try_move_to_virtual_memory(xl_refiner.unet.model)
@@ -290,7 +293,7 @@ def refresh_everything(refiner_model_name, base_model_name, loras):
     refresh_base_model(base_model_name)
     virtual_memory.load_from_virtual_memory(xl_base.unet.model)
 
-    refresh_loras(loras)
+    patch_base(loras, freeu, b1, b2, s1, s2)
     clear_all_caches()
     return
 
@@ -302,7 +305,12 @@ def refresh_everything(refiner_model_name, base_model_name, loras):
         (default_settings['lora_2_model'], default_settings['lora_2_weight']),
         (default_settings['lora_3_model'], default_settings['lora_3_weight']),
         (default_settings['lora_4_model'], default_settings['lora_4_weight']),
-        (default_settings['lora_5_model'], default_settings['lora_5_weight'])]
+        (default_settings['lora_5_model'], default_settings['lora_5_weight'])],
+    freeu=default_settings['freeu'],
+    b1=default_settings['freeu_b1'],
+    b2=default_settings['freeu_b2'],
+    s1=default_settings['freeu_s1'],
+    s2=default_settings['freeu_s2']
 )
 
 expansion = FooocusExpansion()

modules/patch.py

@@ -305,15 +305,24 @@ def blend_latent(a, b, w):
 
 
 def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+    """
+    Apply the model to an input batch.
+    :param x: an [N x C x ...] Tensor of inputs.
+    :param timesteps: a 1-D batch of timesteps.
+    :param context: conditioning plugged in via crossattn
+    :param y: an [N] Tensor of labels, if class-conditional.
+    :return: an [N x C x ...] Tensor of outputs.
+    """
     inpaint_fix = None
     if inpaint_worker.current_task is not None:
         inpaint_fix = inpaint_worker.current_task.inpaint_head_feature
 
     transformer_options["original_shape"] = list(x.shape)
     transformer_options["current_index"] = 0
+    transformer_patches = transformer_options.get("patches", {})
 
     assert (y is not None) == (
-            self.num_classes is not None
+        self.num_classes is not None
     ), "must specify y if and only if the model is class-conditional"
     hs = []
     t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(self.dtype)
@@ -353,6 +362,11 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
             if ctrl is not None:
                 hsp += ctrl
 
+        if "output_block_patch" in transformer_patches:
+            patch = transformer_patches["output_block_patch"]
+            for p in patch:
+                h, hsp = p(h, hsp, transformer_options)
+
         h = torch.cat([h, hsp], dim=1)
         del hsp
         if len(hs) > 0:

modules/settings.py

@@ -65,6 +65,11 @@ def load_settings():
     settings['lora_4_weight'] = modules.path.default_lora_weight
     settings['lora_5_model'] = 'None'
     settings['lora_5_weight'] = modules.path.default_lora_weight
+    settings['freeu'] = False
+    settings['freeu_b1'] = 1.01
+    settings['freeu_b2'] = 1.02
+    settings['freeu_s1'] = 0.99
+    settings['freeu_s2'] = 0.95
 
     if exists('settings.json'):
         with open('settings.json') as settings_file:

readme.md

@@ -189,6 +189,7 @@ Below things are already inside the software, and **users do not need to do anyt
 24. Support for loading models from subfolders (ported from RuinedFooocus).
 25. Support for authentication in --share mode (credentials loaded from auth.json - use auth-example.json as a template).
 26. Support for wildcards (ported from RuinedFooocus - put them in wildcards folder, then try prompts like `__color__ sports car` with different seeds).
+27. Support for [FreeU](https://chenyangsi.top/FreeU/).
 
 ## Thanks
 

settings-example.json

@@ -56,5 +56,10 @@
     "lora_4_model": "None",
     "lora_4_weight": 0.5,
     "lora_5_model": "None",
-    "lora_5_weight": 0.5
+    "lora_5_weight": 0.5,
+    "freeu": false,
+    "freeu_b1": 1.01,
+    "freeu_b2": 1.02,
+    "freeu_s1": 0.99,
+    "freeu_s2": 0.95
 }

update_log_mre.md

@@ -1,3 +1,8 @@
+### 2.0.78.2 MRE
+
+* Added support for FreeU.
+* Updated Comfy.
+
 ### 2.0.78.1 MRE
 
 * Fixed reading paths from paths.json (broken in 2.0.73 MRE).

webui.py

@@ -92,7 +92,7 @@ def metadata_to_ctrls(metadata, ctrls):
     # image_number
     if 'seed' in metadata:
         ctrls[6] = metadata['seed']
-        ctrls[55] = False
+        ctrls[60] = False
     if 'sharpness' in metadata:
         ctrls[7] = metadata['sharpness']
     if 'sampler_name' in metadata:
@@ -210,6 +210,16 @@ def metadata_to_ctrls(metadata, ctrls):
         ctrls[54] = metadata['prompt_expansion']
     elif 'software' in metadata and metadata['software'].startswith('Fooocus 1.'):
         ctrls[54] = False
+    if 'freeu' in metadata:
+        ctrls[55] = metadata['freeu']
+    if 'freeu_b1' in metadata:
+        ctrls[56] = metadata['freeu_b1']
+    if 'freeu_b2' in metadata:
+        ctrls[57] = metadata['freeu_b2']
+    if 'freeu_s1' in metadata:
+        ctrls[58] = metadata['freeu_s1']
+    if 'freeu_s2' in metadata:
+        ctrls[59] = metadata['freeu_s2']
     # seed_random
     return ctrls    
 
@@ -513,6 +523,16 @@ def depth_changed(value):
                 scheduler = gr.Dropdown(label='Scheduler', choices=['karras', 'exponential', 'sgm_uniform', 'simple', 'ddim_uniform'], value=settings['scheduler'])
                 sharpness = gr.Slider(label='Sampling Sharpness', minimum=0.0, maximum=30.0, step=0.01, value=settings['sharpness'])
 
+                freeu_enabled = gr.Checkbox(label='FreeU', value=settings['freeu'])
+                freeu_b1 = gr.Slider(label='Backbone Scaling Factor 1', minimum=0, maximum=2, step=0.01,
+                    value=settings['freeu_b1'], visible=settings['freeu'])
+                freeu_b2 = gr.Slider(label='Backbone Scaling Factor 2', minimum=0, maximum=2, step=0.01,
+                    value=settings['freeu_b2'], visible=settings['freeu'])
+                freeu_s1 = gr.Slider(label='Skip Scaling Factor 1', minimum=0, maximum=2, step=0.01,
+                    value=settings['freeu_s1'], visible=settings['freeu'])
+                freeu_s2 = gr.Slider(label='Skip Scaling Factor 2', minimum=0, maximum=2, step=0.01,
+                    value=settings['freeu_s2'], visible=settings['freeu'])
+
                 def model_refresh_clicked():
                     modules.path.update_all_model_names()
                     results = []
@@ -523,6 +543,12 @@ def model_refresh_clicked():
 
                 model_refresh.click(model_refresh_clicked, [], [base_model, refiner_model] + lora_ctrls, queue=False)
 
+                def freeu_changed(value):
+                    return gr.update(visible=value == True), gr.update(visible=value == True), gr.update(visible=value == True), gr.update(visible=value == True)
+
+                freeu_enabled.change(fn=freeu_changed, inputs=[freeu_enabled], outputs=[freeu_b1, freeu_b2, freeu_s1, freeu_s2])
+
+                freeu_ctrls = [freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2]
 
             with gr.Tab(label='Misc'):
                 output_format = gr.Radio(label='Output Format', choices=['png', 'jpg'], value=settings['output_format'])
@@ -561,7 +587,7 @@ def verify_input(img2img, canny, depth, gallery_in, gallery_rev, gallery_out):
         ]
         ctrls += [base_model, refiner_model, base_clip_skip, refiner_clip_skip] + lora_ctrls
         ctrls += [save_metadata_json, save_metadata_image] + img2img_ctrls + [same_seed_for_all, output_format]
-        ctrls += canny_ctrls + depth_ctrls + [prompt_expansion]
+        ctrls += canny_ctrls + depth_ctrls + [prompt_expansion] + freeu_ctrls
         load_prompt_button.upload(fn=load_prompt_handler, inputs=[load_prompt_button] + ctrls + [seed_random], outputs=ctrls + [seed_random])
         ctrls += [input_image_checkbox, current_tab]
         ctrls += [uov_method, uov_input_image]