Added limited support for non-SDXL models

Author: MoonRide303

fooocus_version.py

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

launch.py

@@ -85,7 +85,9 @@ def prepare_environment():
 
 vae_approx_filenames = [
     ('xlvaeapp.pth',
-     'https://huggingface.co/lllyasviel/misc/resolve/main/xlvaeapp.pth')
+     'https://huggingface.co/lllyasviel/misc/resolve/main/xlvaeapp.pth'),
+    ('taesd_decoder.pth',
+     'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth')
 ]
 
 

modules/async_worker.py

@@ -310,6 +310,13 @@ def handler(task):
             s1=freeu_s1,
             s2=freeu_s2)
 
+        is_sdxl = pipeline.is_base_sdxl()
+        if not is_sdxl:
+            print('WARNING: using non-SDXL base model (supported in limited scope).')
+            control_lora_canny = False
+            control_lora_depth = False
+            revision_mode = False
+
         pipeline.set_clip_skips(base_clip_skip, refiner_clip_skip)
         if revision_mode:
             pipeline.refresh_clip_vision()
@@ -456,10 +463,12 @@ def callback(step, x0, x, total_steps, y):
             input_image = None
             if input_image_path != None:
                 img2img_megapixels = width * height * img2img_scale ** 2 / 2**20
-                if img2img_megapixels < constants.MIN_MEGAPIXELS:
-                    img2img_megapixels = constants.MIN_MEGAPIXELS
-                elif img2img_megapixels > constants.MAX_MEGAPIXELS:
-                    img2img_megapixels = constants.MAX_MEGAPIXELS
+                min_mp = constants.MIN_MEGAPIXELS if is_sdxl else constants.MIN_MEGAPIXELS_SD
+                max_mp = constants.MAX_MEGAPIXELS if is_sdxl else constants.MAX_MEGAPIXELS_SD
+                if img2img_megapixels < min_mp:
+                    img2img_megapixels = min_mp
+                elif img2img_megapixels > max_mp:
+                    img2img_megapixels = max_mp
                 input_image = get_image(input_image_path, img2img_megapixels)
 
             try:

modules/constants.py

@@ -9,6 +9,10 @@
 # exclusive, needed by modules\expansion.py -> transformers\trainer_utils.py -> np.random.seed()
 SEED_LIMIT_NUMPY = 2**32
 
-# min - native SDXL resolution, max - determined by SDXL context size (2048)
+# min - native SDXL resolution (1024x1024), max - determined by SDXL context size (2048)
 MIN_MEGAPIXELS = 1.0
 MAX_MEGAPIXELS = 4.0
+
+# min - native SD 1.5 resolution (512x512), max - determined by SD 2.x context size (1024)
+MIN_MEGAPIXELS_SD = 0.25
+MAX_MEGAPIXELS_SD = 1.0

modules/core.py

@@ -19,7 +19,7 @@
 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.model_base import SDXL, SDXLRefiner
 from comfy.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
 from modules.samplers_advanced import KSamplerBasic, KSamplerWithRefiner
 from modules.path import embeddings_path
@@ -236,14 +236,15 @@ def forward(self, x):
 
 
 VAE_approx_model = None
+taesd = None
 
 
 @torch.no_grad()
 @torch.inference_mode()
-def get_previewer(device, latent_format):
-    global VAE_approx_model
+def get_previewer(device, latent_format, is_sdxl=True):
+    global VAE_approx_model, taesd
 
-    if VAE_approx_model is None:
+    if VAE_approx_model is None and is_sdxl:
         from modules.path import vae_approx_path
         vae_approx_filename = os.path.join(vae_approx_path, 'xlvaeapp.pth')
         sd = torch.load(vae_approx_filename, map_location='cpu')
@@ -271,8 +272,26 @@ def preview_function(x0, step, total_steps):
             x_sample = x_sample.cpu().numpy().clip(0, 255).astype(np.uint8)
             return x_sample
 
-    return preview_function
+    if taesd is None and not is_sdxl:
+        from latent_preview import TAESD, TAESDPreviewerImpl
+        taesd_decoder_path = os.path.abspath(os.path.realpath(os.path.join("models", "vae_approx", latent_format.taesd_decoder_name)))
 
+        if not os.path.exists(taesd_decoder_path):
+            print(f"Warning: TAESD previews enabled, but could not find {taesd_decoder_path}")
+            return None
+
+        taesd = TAESD(None, taesd_decoder_path).to(device)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def preview_function_sd(x0, step, total_steps):
+        with torch.no_grad():
+            x_sample = taesd.decoder(torch.nn.functional.avg_pool2d(x0, kernel_size=(2, 2))).detach() * 255.0
+            x_sample = einops.rearrange(x_sample, 'b c h w -> b h w c')
+            x_sample = x_sample.cpu().numpy().clip(0, 255).astype(np.uint8)
+            return x_sample[0]
+
+    return preview_function if is_sdxl else preview_function_sd
 
 @torch.no_grad()
 @torch.inference_mode()
@@ -299,7 +318,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
-    previewer = get_previewer(device, model.model.latent_format)
+    previewer = get_previewer(device, model.model.latent_format, isinstance(model.model, SDXL))
 
     pbar = comfy.utils.ProgressBar(steps)
 
@@ -372,7 +391,7 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
-    previewer = get_previewer(device, model.model.latent_format)
+    previewer = get_previewer(device, model.model.latent_format, isinstance(model.model, SDXL))
 
     pbar = comfy.utils.ProgressBar(steps)
 

modules/default_pipeline.py

@@ -7,7 +7,7 @@
 import modules.virtual_memory as virtual_memory
 import comfy.model_management
 
-from comfy.model_base import SDXL, SDXLRefiner
+from comfy.model_base import BaseModel, SDXL, SDXLRefiner
 from modules.settings import default_settings
 from modules.patch import set_comfy_adm_encoding, set_fooocus_adm_encoding, cfg_patched, patched_model_function
 from modules.expansion import FooocusExpansion
@@ -48,8 +48,8 @@ def refresh_base_model(name):
         xl_base = None
 
     xl_base = core.load_model(filename)
-    if not isinstance(xl_base.unet.model, SDXL):
-        print('Model not supported. Fooocus only support SDXL model as the base model.')
+    if not isinstance(xl_base.unet.model, BaseModel):
+        print(f'Model not supported: {name}, using default base model instead.')
         xl_base = None
         xl_base_hash = ''
         refresh_base_model(modules.path.default_base_model_name)
@@ -58,13 +58,21 @@ def refresh_base_model(name):
         xl_base_patched_hash = ''
         return
 
+    if not isinstance(xl_base.unet.model, SDXL):
+        print('WARNING: loading non-SDXL base model.')
+
     xl_base_hash = model_hash
     xl_base_patched = xl_base
     xl_base_patched_hash = ''
     print(f'Base model loaded: {model_hash}')
     return
 
 
+def is_base_sdxl():
+    assert xl_base is not None
+    return isinstance(xl_base.unet.model, SDXL)
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def refresh_refiner_model(name):
@@ -369,7 +377,7 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
         positive_conditions, negative_conditions = core.apply_controlnet(positive_conditions, negative_conditions,
             controlnet_depth, input_image, depth_strength, depth_start, depth_stop)
 
-    if xl_refiner is not None:
+    if xl_refiner is not None and is_base_sdxl():
         positive_conditions_refiner = positive_cond[1]
         negative_conditions_refiner = negative_cond[1]
 

readme.md

@@ -190,6 +190,7 @@ Below things are already inside the software, and **users do not need to do anyt
 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/).
+28. Limited support for non-SDXL models (no refiner, Control-LoRAs, Revision, inpainting, outpainting).
 
 ## Thanks
 

update_log_mre.md

@@ -1,3 +1,7 @@
+### 2.0.78.3 MRE
+
+* Added limited support for non-SDXL models (no refiner, Control-LoRAs, Revision, inpainting, outpainting).
+
 ### 2.0.78.2 MRE
 
 * Added support for FreeU.