From 51e8567eec51a72de792d9135e7884e9664fd615 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 8 Sep 2024 10:48:22 -0400 Subject: [PATCH] flux hires and refiner workflows --- CHANGELOG.md | 10 +++++++--- modules/processing_args.py | 13 +++++++++++++ modules/processing_callbacks.py | 17 +++++++++++++---- modules/processing_diffusers.py | 22 ++++++++++++---------- modules/processing_helpers.py | 6 +++--- modules/processing_vae.py | 8 ++++++-- modules/sd_samplers.py | 2 +- 7 files changed, 55 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ad29dee3..1b2447b0f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/modules/processing_args.py b/modules/processing_args.py index d08783256..ef7e9f30f 100644 --- a/modules/processing_args.py +++ b/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') diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index a9c7ab5e8..52142d44b 100644 --- a/modules/processing_callbacks.py +++ b/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 diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index db9e3425c..a675d799f 100644 --- a/modules/processing_diffusers.py +++ b/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: diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 6dfbdc535..31e8c3ba4 100644 --- a/modules/processing_helpers.py +++ b/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) diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 1e7900ff0..83d557b9b 100644 --- a/modules/processing_vae.py +++ b/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) diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 61e17df8f..b7c90e603 100644 --- a/modules/sd_samplers.py +++ b/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)