flux hires and refiner workflows

CHANGELOG.md

@@ -1,13 +1,14 @@
 # Change Log for SD.Next
 
-## Update for 2024-09-07
+## Update for 2024-09-08
 
-### Highlights for 2024-09-07
+### Highlights for 2024-09-08
 
 Major refactor of [FLUX.1](https://blackforestlabs.ai/announcing-black-forest-labs/) support:  
 - Full **ControlNet** support, better **LoRA** support, full **prompt attention** implementation  
 - Faster execution, more flexible loading, additional quantization options, and more...  
 - Added **image-to-image**, **inpaint**, **outpaint**, **hires** modes  
+- Added workflow where FLUX can be used as **refiner** for other models  
 - Since both *Optimum-Quanto* and *BitsAndBytes* libraries are limited in their platform support matrix,  
   try enabling **NNCF** for quantization/compression on-the-fly!  
 
@@ -23,7 +24,7 @@ And few video related goodies...
 
 Plus tons of minor items and fixes - see [changelog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) for details!  
 
-### Details for 2024-09-07
+### Details for 2024-09-08
 
 **Major refactor of FLUX.1 support:**
 - allow configuration of individual FLUX.1 model components: *transformer, text-encoder, vae*  
@@ -46,6 +47,9 @@ Plus tons of minor items and fixes - see [changelog](https://github.com/vladmand
   not recommended due to massive duplication of components, but added due to popular demand  
   each such model is 20-32GB in size vs ~11GB for typical unet fine-tune  
 - improve logging, warn when attempting to load unet as base model  
+- **refiner** support  
+  FLUX.1 can be used as refiner for other models such as sd/sdxl  
+  simply load sd/sdxl model as base and flux model as refiner and use as usual refiner workflow  
 - **img2img**, **inpaint** and **outpaint** support  
   *note* flux may require higher denoising strength than typical sd/sdxl models  
   *note*: img2img is not yet supported with controlnet  

modules/processing_args.py

@@ -34,6 +34,10 @@ def task_specific_kwargs(p, model):
             'image': p.init_images,
             'strength': p.denoising_strength,
         }
+        if model.__class__.__name__ == 'FluxImg2ImgPipeline': # needs explicit width/height
+            p.width = 8 * math.ceil(p.init_images[0].width / 8)
+            p.height = 8 * math.ceil(p.init_images[0].height / 8)
+            task_args['width'], task_args['height'] = p.width, p.height
     elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0:
         p.ops.append('instruct')
         task_args = {
@@ -229,6 +233,15 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
             args['cross_attention_kwargs'] = {}
         args['cross_attention_kwargs'][k] = v
 
+    # handle missing resolution
+    if args.get('image', None) is not None and ('width' not in args or 'height' not in args):
+        if isinstance(args['image'], torch.Tensor) or isinstance(args['image'], np.ndarray):
+            args['width'] = 8 * args['image'].shape[-1]
+            args['height'] = 8 * args['image'].shape[-2]
+        else:
+            args['width'] = 8 * math.ceil(args['image'][0].width / 8)
+            args['height'] = 8 * math.ceil(args['image'][0].height / 8)
+
     # handle implicit controlnet
     if 'control_image' in possible and 'control_image' not in args and 'image' in args:
         debug('Diffusers: set control image')

modules/processing_callbacks.py

@@ -82,10 +82,19 @@ 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]
-    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) # pylint: disable=protected-access
-    else:
-        shared.state.current_latent = kwargs['latents']
+    try:
+        if hasattr(pipe, "_unpack_latents") and hasattr(pipe, "vae_scale_factor"): # FLUX
+            if p.hr_upscaler is not None and p.hr_upscaler != 'None':
+                width = max(getattr(p, 'width', 0), getattr(p, 'hr_upscale_to_x', 0))
+                height = max(getattr(p, 'height', 0), getattr(p, 'hr_upscale_to_y', 0))
+            else:
+                width = getattr(p, 'width', 0)
+                height = getattr(p, 'height', 0)
+            shared.state.current_latent = pipe._unpack_latents(kwargs['latents'], height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
+        else:
+            shared.state.current_latent = kwargs['latents']
+    except Exception as e:
+        shared.log.error(f'Callback: {e}')
     if shared.cmd_opts.profile and shared.profiler is not None:
         shared.profiler.step()
     return kwargs

modules/processing_diffusers.py

@@ -181,7 +181,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             shared.log.info(f'HiRes: class={shared.sd_model.__class__.__name__} sampler="{p.hr_sampler_name}"')
             if p.is_control and hasattr(p, 'task_args') and p.task_args.get('image', None) is not None:
                 if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
-                    output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality, output_type='pil') # controlnet cannnot deal with latent input
+                    output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input
                     p.task_args['image'] = output.images # replace so hires uses new output
             sd_models.move_model(shared.sd_model, devices.device)
             orig_denoise = p.denoising_strength
@@ -246,8 +246,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             image = output.images[i]
             noise_level = round(350 * p.denoising_strength)
             output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np'
-            if shared.sd_refiner.__class__.__name__ == 'StableDiffusionUpscalePipeline':
-                image = processing_vae.vae_decode(latents=image, model=shared.sd_model, full_quality=p.full_quality, output_type='pil')
+            if 'Upscale' in shared.sd_refiner.__class__.__name__ or 'Flux in shared.sd_refiner.__class__.__name__':
+                image = processing_vae.vae_decode(latents=image, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height)
                 p.extra_generation_params['Noise level'] = noise_level
                 output_type = 'np'
             if hasattr(p, 'task_args') and p.task_args.get('image', None) is not None and output is not None: # replace input with output so it can be used by hires/refine
@@ -284,7 +284,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 shared.log.info(e)
 
             if not shared.state.interrupted and not shared.state.skipped:
-                refiner_images = processing_vae.vae_decode(latents=refiner_output.images, model=shared.sd_refiner, full_quality=True)
+                refiner_images = processing_vae.vae_decode(latents=refiner_output.images, model=shared.sd_refiner, full_quality=True, width=max(p.width, p.hr_upscale_to_x), height=max(p.height, p.hr_upscale_to_y))
                 for refiner_image in refiner_images:
                     results.append(refiner_image)
 
@@ -303,12 +303,14 @@ 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) # pylint: disable=protected-access
-            if torch.is_tensor(output.images) and len(output.images) > 0 and any(s >= 512 for s in output.images.shape):
-                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)
+            if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
+                if p.hr_upscaler is not None and p.hr_upscaler != 'None':
+                    width = max(getattr(p, 'width', 0), getattr(p, 'hr_upscale_to_x', 0))
+                    height = max(getattr(p, 'height', 0), getattr(p, 'hr_upscale_to_y', 0))
+                else:
+                    width = getattr(p, 'width', 0)
+                    height = getattr(p, 'height', 0)
+                results = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality, width=width, height=height)
             elif hasattr(output, 'images'):
                 results = output.images
             else:

modules/processing_helpers.py

@@ -383,13 +383,13 @@ def resize_init_images(p):
 def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler else latent
     if not torch.is_tensor(latents):
         shared.log.warning('Hires: input is not tensor')
-        first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil')
+        first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height)
         return first_pass_images
     latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None)
     # shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}')
     if latent_upscaler is not None:
         return torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=latent_upscaler["mode"], antialias=latent_upscaler["antialias"])
-    first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil')
+    first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height)
     if p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0 and hasattr(p, 'init_hr'):
         shared.log.error('Hires: missing upscaling dimensions')
         return first_pass_images
@@ -531,7 +531,7 @@ def save_intermediate(p, latents, suffix):
     for i in range(len(latents)):
         from modules.processing import create_infotext
         info=create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i)
-        decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, output_type='pil', full_quality=p.full_quality)
+        decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, output_type='pil', full_quality=p.full_quality, width=p.width, height=p.height)
         for j in range(len(decoded)):
             images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
 

modules/processing_vae.py

@@ -115,7 +115,7 @@ def taesd_vae_encode(image):
     return encoded
 
 
-def vae_decode(latents, model, output_type='np', full_quality=True):
+def vae_decode(latents, model, output_type='np', full_quality=True, width=None, height=None):
     t0 = time.time()
     prev_job = shared.state.job
     shared.state.job = 'VAE'
@@ -129,11 +129,15 @@ def vae_decode(latents, model, output_type='np', full_quality=True):
     if not hasattr(model, 'vae'):
         shared.log.error('VAE not found in model')
         return []
+    if hasattr(model, "_unpack_latents") and hasattr(model, "vae_scale_factor") and width is not None and height is not None: # FLUX
+        latents = model._unpack_latents(latents, height, width, model.vae_scale_factor) # pylint: disable=protected-access
     if len(latents.shape) == 3: # lost a batch dim in hires
         latents = latents.unsqueeze(0)
     if latents.shape[0] == 4 and latents.shape[1] != 4: # likely animatediff latent
         latents = latents.permute(1, 0, 2, 3)
-    if full_quality:
+    if any(s >= 512 for s in latents.shape):
+        imgs = latents.float().cpu().numpy()
+    elif full_quality and hasattr(shared.sd_model, "vae"):
         decoded = full_vae_decode(latents=latents, model=shared.sd_model)
     else:
         decoded = taesd_vae_decode(latents=latents)

modules/sd_samplers.py

@@ -69,7 +69,7 @@ def create_sampler(name, model):
         return sampler
     elif shared.native:
         sampler = config.constructor(model)
-        if shared.sd_model_type == 'f1':
+        if 'Flux' in model.__class__.__name__:
             if 'base_image_seq_len' not in sampler.sampler.config or 'max_image_seq_len' not in sampler.sampler.config or 'base_shift' not in sampler.sampler.config or 'max_shift' not in sampler.sampler.config:
                 shared.log.warning(f'FLUX: sampler="{name}" unsupported')
                 # sampler.sampler.register_to_config(base_image_seq_len=256, max_image_seq_len=4096, base_shift=0.5, max_shift=1.15)