FLUX support

modules/processing_callbacks.py

@@ -82,7 +82,10 @@ def diffusers_callback(pipe, step: int, timestep: int, kwargs: dict):
             pipe._guidance_scale = 0.0  # pylint: disable=protected-access
             for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"} & set(kwargs):
                 kwargs[key] = kwargs[key].chunk(2)[-1]
-    shared.state.current_latent = kwargs['latents']
+    if hasattr(pipe, "_unpack_latents") and hasattr(pipe, "vae_scale_factor"): # FLUX
+        shared.state.current_latent = pipe._unpack_latents(kwargs['latents'], p.height, p.width, pipe.vae_scale_factor)
+    else:
+        shared.state.current_latent = kwargs['latents']
     if shared.cmd_opts.profile and shared.profiler is not None:
         shared.profiler.step()
     return kwargs

modules/processing_diffusers.py

@@ -309,8 +309,10 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             if not hasattr(output, 'images') and hasattr(output, 'frames'):
                 shared.log.debug(f'Generated: frames={len(output.frames[0])}')
                 output.images = output.frames[0]
+            if hasattr(shared.sd_model, "_unpack_latents") and hasattr(shared.sd_model, "vae_scale_factor"): # FLUX
+                output.images = shared.sd_model._unpack_latents(output.images, p.height, p.width, shared.sd_model.vae_scale_factor)
             if torch.is_tensor(output.images) and len(output.images) > 0 and any(s >= 512 for s in output.images.shape):
-                results = output.images.cpu().numpy()
+                results = output.images.float().cpu().numpy()
             elif hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
                 results = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality)
             elif hasattr(output, 'images'):

modules/sd_models_compile.py

@@ -145,6 +145,7 @@ def nncf_compress_model(model):
 def nncf_compress_weights(sd_model):
     try:
         t0 = time.time()
+        shared.log.info(f"NNCF Compress Weights: {shared.opts.nncf_compress_weights}")
         from installer import install
         install('nncf==2.7.0', quiet=True)
 
@@ -173,6 +174,7 @@ def optimum_quanto_model(model, weights=None):
 def optimum_quanto_weights(sd_model):
     try:
         t0 = time.time()
+        shared.log.info(f"Optimum Quanto Weights: {shared.opts.optimum_quanto_weights}")
         from installer import install
         install('optimum-quanto', quiet=True)