From e2d13a4dfbc6931a7f528d9f5e3754cb04d4f448 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 2 Sep 2024 11:54:18 -0400 Subject: [PATCH] flux controlnet support --- CHANGELOG.md | 37 ++++++++++++++++++++++--------------- javascript/sdnext.css | 3 ++- modules/control/run.py | 24 +++++++++++++----------- modules/control/unit.py | 11 +++++++++++ modules/modeldata.py | 2 +- modules/processing_class.py | 1 + modules/sd_models.py | 31 ++++++++++++++++++------------- modules/ui_control.py | 2 ++ 8 files changed, 70 insertions(+), 41 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e84b92245..453a2ee4f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,14 @@ # Change Log for SD.Next -## Update for 2024-09-01 +## Update for 2024-09-02 -Major refactor of FLUX.1 support: +### Highlights + +Major refactor of FLUX.1 support: faster, more flexible loading, full ControlNet support, better LoRA support, full prompt attention support, additional quantization options, and more... + +### Details + +**Major refactor of FLUX.1 support:** - allow configuration of individual FLUX.1 model components: *transformer, text-encoder, vae* model load will load selected components first and then initialize model using pre-loaded components components that were not pre-loaded will be downloaded and initialized as needed @@ -13,18 +19,24 @@ Major refactor of FLUX.1 support: - transformer/unet is list of manually downloaded safetensors - vae is list of manually downloaded safetensors - text-encoder is list of predefined and manually downloaded text-encoders -- model support loading all-in-one safetensors (*1) +- **controlnet** support: (*1) + support for **InstantX/Shakker-Labs** models including [Union-Pro](InstantX/FLUX.1-dev-Controlnet-Union) + note that flux controlnet models are large, up to 6.6GB on top of already large base model! + as such, you may need to use offloading:sequential which is not as fast, but uses far less memory + when using union model, you must also select control mode in the control unit + flux does not yet support *img2img* so to use controlnet, you need to set contronet input via control unit override +- model support loading **all-in-one** safetensors (*1) 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 -- transformer/unet support fp8/fp4 quantization -- vae support fp16 (*1) -- lora support additional training tools (*1) +- transformer/unet support *fp8/fp4* quantization +- vae support *fp16* (*1) +- **lora** support additional training tools (*1) - support fuse-qkv projections (*1) can speed up generate enable via *settings -> compute -> fused projections* -Other improvements: +**Other improvements:** - taesd configurable number of layers can be used to speed-up taesd decoding by reducing number of ops e.g. if generating 1024px image, reducing layers by 1 will result in preview being 512px @@ -39,18 +51,13 @@ Other improvements: - t5 support manually downloaded models applies to all models that use t5 transformer -Fixes: +**Fixes:** - fix handling of model configs if offline config is not available - fix vae decode in backend original Work-in-progress: -- flux controlnet support: (*1) - - instanx models - - shakker-labs models - - TBD: add controlnet_mode for union models - - TBD: validate control_image vs input_type - - TBD: not enough values to unpack - - TBD: xlabs models +- TBD: flux controlnet XLabs-AI models +- TBD: flux img2img/inpaint *notes*: - (*1) requires `diffusers==0.31.0.dev0` diff --git a/javascript/sdnext.css b/javascript/sdnext.css index 6f6cfc200..d2b59771b 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -285,11 +285,12 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt #control_script_container { display: block; margin-top: 1em; border-width: 2px 0 0 0; border-style: solid; border-color: var(--highlight-color); } .control-button { min-height: 42px; max-height: 42px; line-height: 1em; } .control-tabs > .tab-nav { margin-bottom: 0; margin-top: 0; } -.control-unit { max-width: 1200px; padding: 0 !important; margin-top: -10px !important; } +.control-unit { padding: 0 !important; margin-top: -10px !important; } .control-unit > .label-wrap { margin-bottom: 0 !important; } .control-settings { border-width: var(--block-border-width) !important; border-top: var(--button-primary-border-color) !important; border-style: solid !important; margin-top: 1em !important; } .processor-settings { padding: 0 !important; max-width: 300px; } .processor-group > div { flex-flow: wrap;gap: 1em; } +.control-unit .gradio-button.tool { align-self: baseline; margin-top: 2rem; } /* main info */ .main-info { font-weight: var(--section-header-text-weight); color: var(--body-text-color-subdued); padding: 1em !important; margin-top: 2em !important; line-height: var(--line-lg) !important; } diff --git a/modules/control/run.py b/modules/control/run.py index f978dbf36..a69b30198 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -234,7 +234,8 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini active_start.append(float(u.start)) active_end.append(float(u.end)) p.guess_mode = u.guess - shared.log.debug(f'Control ControlNet unit: i={num_units} process={u.process.processor_id} model={u.controlnet.model_id} strength={u.strength} guess={u.guess} start={u.start} end={u.end}') + p.control_mode = u.mode + shared.log.debug(f'Control ControlNet unit: i={num_units} process={u.process.processor_id} model={u.controlnet.model_id} strength={u.strength} guess={u.guess} start={u.start} end={u.end} mode={u.mode}') elif unit_type == 'xs' and u.controlnet.model is not None: active_process.append(u.process) active_model.append(u.controlnet) @@ -388,7 +389,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini video = cv2.VideoCapture(inputs) if not video.isOpened(): if is_generator: - yield terminate(f'Control: video open failed: path={inputs}') + yield terminate(f'Video open failed: path={inputs}') return [], '', '', 'Error: video open failed' frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) fps = int(video.get(cv2.CAP_PROP_FPS)) @@ -401,7 +402,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini shared.log.debug(f'Control: input video: path={inputs} frames={frames} fps={fps} size={w}x{h} codec={codec}') except Exception as e: if is_generator: - yield terminate(f'Control: video open failed: path={inputs} {e}') + yield terminate(f'Video open failed: path={inputs} {e}') return [], '', '', 'Error: video open failed' while status: @@ -419,7 +420,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini if shared.state.interrupted: shared.state.interrupted = False if is_generator: - yield terminate('Control interrupted') + yield terminate('Interrupted') return [], '', '', 'Interrupted' # get input if isinstance(input_image, str): @@ -508,7 +509,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini pass if any(img is None for img in processed_images): if is_generator: - yield terminate('Control: attempting process but output is none') + yield terminate('Attempting process but output is none') return [], '', '', 'Error: output is none' if len(processed_images) > 1 and len(active_process) != len(active_model): processed_image = [np.array(i) for i in processed_images] @@ -527,7 +528,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini p.init_images = processed_images elif isinstance(selected_models, list) and len(processed_images) != len(selected_models): if is_generator: - yield terminate(f'Control: number of inputs does not match: input={len(processed_images)} models={len(selected_models)}') + yield terminate(f'Number of inputs does not match: input={len(processed_images)} models={len(selected_models)}') return [], '', '', 'Error: number of inputs does not match' elif selected_models is not None: p.init_images = processed_image @@ -542,14 +543,13 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini debug(f'Control: process=None image={p.ref_image}') if p.ref_image is None: if is_generator: - yield terminate('Control: attempting reference mode but image is none') + yield terminate('Attempting reference mode but image is none') return [], '', '', 'Reference mode without image' elif unit_type == 'controlnet' and has_models: if input_type == 0: # Control only if shared.sd_model_type == 'f1': - if is_generator: - yield terminate('Control: Flux control invalid input type') - return [], '', '', 'Flux control invalid input type' + p.task_args['control_image'] = p.init_images # TODO flux controlnet mandates this + p.task_args['strength'] = p.denoising_strength elif input_type == 1: # Init image same as control p.task_args['control_image'] = p.init_images # switch image and control_image p.task_args['strength'] = p.denoising_strength @@ -602,6 +602,8 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini if unit_type == 'lite': p.init_image = [input_image] instance.apply(selected_models, processed_image, control_conditioning) + if p.control_mode is not None: + p.task_args['control_mode'] = p.control_mode if hasattr(p, 'init_images') and p.init_images is None: # delete empty del p.init_images @@ -609,7 +611,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini if has_models: if unit_type in ['controlnet', 't2i adapter', 'lite', 'xs'] and p.task_args.get('image', None) is None and getattr(p, 'init_images', None) is None: if is_generator: - yield terminate(f'Control: mode={p.extra_generation_params.get("Control mode", None)} input image is none') + yield terminate(f'Mode={p.extra_generation_params.get("Control mode", None)} input image is none') return [], '', '', 'Error: Input image is none' # resize mask diff --git a/modules/control/unit.py b/modules/control/unit.py index b2b102b9b..51101493d 100644 --- a/modules/control/unit.py +++ b/modules/control/unit.py @@ -36,6 +36,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c image_preview = None, control_start = None, control_end = None, + control_mode = None, result_txt = None, extra_controls: list = [], ): @@ -46,6 +47,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c self.end = end or 1 self.start = min(self.start, self.end) self.end = max(self.start, self.end) + self.mode = None # processor always exists, adapter and controlnet are optional self.process: processors.Processor = processors.Processor() self.adapter: t2iadapter.Adapter = None @@ -83,6 +85,12 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c self.start = min(start, end) self.end = max(start, end) + def control_mode_change(mode): + self.mode = mode - 1 if mode > 0 else None + + def control_mode_show(model_id): + return gr.update(visible='union' in model_id.lower()) + def adapter_extra(c1): self.factor = c1 @@ -156,6 +164,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c self.controlnet.load(model_id) else: model_id.change(fn=self.controlnet.load, inputs=[model_id], outputs=[result_txt], show_progress=True) + model_id.change(fn=control_mode_show, inputs=[model_id], outputs=[control_mode], show_progress=False) if extra_controls is not None and len(extra_controls) > 0: extra_controls[0].change(fn=controlnet_extra, inputs=extra_controls) elif self.type == 'xs': @@ -202,3 +211,5 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c if control_start is not None and control_end is not None: control_start.change(fn=control_change, inputs=[control_start, control_end]) control_end.change(fn=control_change, inputs=[control_start, control_end]) + if control_mode is not None: + control_mode.change(fn=control_mode_change, inputs=[control_mode]) diff --git a/modules/modeldata.py b/modules/modeldata.py index f002ca7bd..aedbf99e2 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -103,7 +103,7 @@ class Shared(sys.modules[__name__].__class__): model_type = 'sc' elif "AuraFlow" in self.sd_model.__class__.__name__: model_type = 'auraflow' - elif "FluxPipeline" in self.sd_model.__class__.__name__: + elif "FluxPipeline" in self.sd_model.__class__.__name__ or "FluxControlNetPipeline" in self.sd_model.__class__.__name__: model_type = 'f1' else: model_type = self.sd_model.__class__.__name__ diff --git a/modules/processing_class.py b/modules/processing_class.py index ad083ec5a..830cdc618 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -456,6 +456,7 @@ class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img): self.controlnet_conditioning_scale = None self.control_guidance_start = None self.control_guidance_end = None + self.control_mode = None self.reference_attn = None self.reference_adain = None self.attention_auto_machine_weight = None diff --git a/modules/sd_models.py b/modules/sd_models.py index 19fdee91f..bc819e53f 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -689,20 +689,21 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): if shared.opts.no_half_vae: devices.dtype_vae = torch.float32 sd_model.vae.to(devices.dtype_vae) - shared.log.debug(f'Setting {op} VAE: no-half') + shared.log.debug(f'Setting {op} VAE: no-half=True') if hasattr(sd_model, "enable_vae_slicing"): if shared.opts.diffusers_vae_slicing: - shared.log.debug(f'Setting {op}: enable VAE slicing') + shared.log.debug(f'Setting {op}: slicing=True') sd_model.enable_vae_slicing() else: sd_model.disable_vae_slicing() if hasattr(sd_model, "enable_vae_tiling"): if shared.opts.diffusers_vae_tiling: - shared.log.debug(f'Setting {op}: enable VAE tiling') + shared.log.debug(f'Setting {op}: tiling=True') sd_model.enable_vae_tiling() else: sd_model.disable_vae_tiling() if hasattr(sd_model, "vqvae"): + shared.log.debug(f'Setting {op} VQVAE: upcast=True') sd_model.vqvae.to(torch.float32) # vqvae is producing nans in fp16 set_diffusers_attention(sd_model) @@ -710,13 +711,13 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'fuse_qkv_projections'): try: sd_model.fuse_qkv_projections() - shared.log.debug(f'Setting {op}: enable fused projections') + shared.log.debug(f'Setting {op}: fused-qkv=True') except Exception as e: shared.log.error(f'Error enabling fused projections: {e}') if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'fuse_qkv_projections'): try: sd_model.transformer.fuse_qkv_projections() - shared.log.debug(f'Setting {op}: enable fused projections') + shared.log.debug(f'Setting {op}: fused-qkv=True') except Exception as e: shared.log.error(f'Error enabling fused projections: {e}') if shared.opts.diffusers_eval: @@ -730,7 +731,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): sd_model = sd_models_compile.dynamic_quantization(sd_model) if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'): - shared.log.debug(f'Setting {op}: enable channels last') + shared.log.debug(f'Setting {op}: channels-last=True') sd_model.unet.to(memory_format=torch.channels_last) if offload: @@ -743,13 +744,12 @@ def set_diffuser_offload(sd_model, op: str = 'model'): if sd_model is None: shared.log.warning(f'{op} is not loaded') return - shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if hasattr(sd_model, "enable_model_cpu_offload"): if shared.opts.diffusers_offload_mode == "model": try: - shared.log.debug(f'Setting {op}: enable model CPU offload') + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False shared.opts.diffusers_move_unet = False @@ -765,7 +765,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'): if hasattr(sd_model, "enable_sequential_cpu_offload"): if shared.opts.diffusers_offload_mode == "sequential": try: - shared.log.debug(f'Setting {op}: enable sequential CPU offload') + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False shared.opts.diffusers_move_unet = False @@ -785,6 +785,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'): shared.log.error(f'Model offload error: mode={shared.opts.diffusers_offload_mode} {e}') if shared.opts.diffusers_offload_mode == "balanced": try: + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') sd_model = apply_balanced_offload(sd_model) except Exception as e: shared.log.error(f'Model offload error: mode={shared.opts.diffusers_offload_mode} {e}') @@ -840,8 +841,6 @@ def apply_balanced_offload(sd_model): shared.log.error(f'Balanced offload: module={module_name} {e}') devices.torch_gc(fast=True) - if not shared.native: - return apply_balanced_offload_to_module(sd_model) if hasattr(sd_model, "prior_pipe"): apply_balanced_offload_to_module(sd_model.prior_pipe) @@ -861,6 +860,7 @@ def normalize_device(device): return torch.device(str(device) + ":0") return torch.device(device) + def move_model(model, device=None, force=False): if model is None or device is None: return @@ -1558,11 +1558,16 @@ def set_diffusers_attention(pipe): for module in modules: if module.__class__.__name__ in ['SD3Transformer2DModel']: module.set_attn_processor(p.JointAttnProcessor2_0()) - elif module.__class__.__name__ in ['HunyuanDiT2DModel', 'FluxTransformer2DModel']: - pass + elif module.__class__.__name__ in ['FluxTransformer2DModel']: + module.set_attn_processor(p.FluxAttnProcessor2_0()) + elif module.__class__.__name__ in ['HunyuanDiT2DModel']: + module.set_attn_processor(p.HunyuanAttnProcessor2_0()) else: module.set_attn_processor(attention) + if 'ControlNet' in pipe.__class__.__name__: # do not replace attention in ControlNet pipelines + return + shared.log.debug(f"Setting model: attention={shared.opts.cross_attention_optimization}") if shared.opts.cross_attention_optimization == "Disabled": pass # do nothing elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers diff --git a/modules/ui_control.py b/modules/ui_control.py index c4aa25bd3..336653396 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -204,6 +204,7 @@ def create_ui(_blocks: gr.Blocks=None): model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=2.0, step=0.01, value=1.0-i/10) control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0) control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0) + control_mode = gr.Dropdown(label="Mode", choices=['', 'Canny', 'Tile', 'Depth', 'Blur', 'Pose', 'Gray', 'LQ'], value=0, type='index', visible=False) reset_btn = ui_components.ToolButton(value=ui_symbols.reset) image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool']) image_reuse= ui_components.ToolButton(value=ui_symbols.reuse) @@ -226,6 +227,7 @@ def create_ui(_blocks: gr.Blocks=None): image_preview = image_preview, control_start = control_start, control_end = control_end, + control_mode = control_mode, extra_controls = extra_controls, ) )