From 6a5e253ecf64867c9477f863857cdb94458f6a18 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 26 Mar 2025 10:50:09 -0400 Subject: [PATCH] add infiniteyou Signed-off-by: Vladimir Mandic --- .pylintrc | 2 + .ruff.toml | 2 + CHANGELOG.md | 13 +- TODO.md | 2 - modules/infiniteyou/__init__.py | 2 + .../infiniteyou/pipeline_flux_infusenet.py | 612 ++++++++++++++++++ modules/infiniteyou/pipeline_infu_flux.py | 325 ++++++++++ modules/infiniteyou/resampler.py | 121 ++++ modules/lora/networks.py | 2 +- modules/processing_args.py | 10 +- scripts/infiniteyou_ext.py | 121 ++++ 11 files changed, 1200 insertions(+), 12 deletions(-) create mode 100644 modules/infiniteyou/__init__.py create mode 100644 modules/infiniteyou/pipeline_flux_infusenet.py create mode 100644 modules/infiniteyou/pipeline_infu_flux.py create mode 100644 modules/infiniteyou/resampler.py create mode 100644 scripts/infiniteyou_ext.py diff --git a/.pylintrc b/.pylintrc index b472d26a0..bbe29b197 100644 --- a/.pylintrc +++ b/.pylintrc @@ -38,6 +38,8 @@ ignore-paths=/usr/lib/.*$, modules/todo, modules/unipc, modules/xadapter, + modules/infiniteyou, + modules/flash_attn_triton_amd, repositories, extensions-builtin/Lora, extensions-builtin/sd-webui-agent-scheduler, diff --git a/.ruff.toml b/.ruff.toml index 0fc9de8b3..48f2e9026 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -33,6 +33,8 @@ exclude = [ "modules/todo", "modules/unipc", "modules/xadapter", + "modules/infiniteyou", + "modules/flash_attn_triton_amd", "repositories", "extensions-builtin/Lora", "extensions-builtin/sd-extension-chainner/nodes", diff --git a/CHANGELOG.md b/CHANGELOG.md index ffbc2922a..072af4c91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,8 @@ # Change Log for SD.Next -## Update for 2025-03-25 +## Update for 2025-03-26 -### Highlights for 2025-03-25 +### Highlights for 2025-03-26 Brand new Video processing module with support for all latest models: **WAN21, Hunyuan, LTX, Cog, Allegro, Mochi1, Latte1** in both T2V and I2V workflows And combined with *on-the-fly quantization*, support for *Local/Tiny/Remote* VAE, acceleration modules such as *FasterCache or PAB* and more! @@ -11,7 +11,7 @@ Plus support for **CogView-4**, **SANA 1.5**, new CLiP models, improvements to r **Flash Attention 2** and Sage Attention is now available on ZLUDA backend! -### Details for 2025-03-25 +### Details for 2025-03-26 - **Video tab** - see [Video Wiki](https://github.com/vladmandic/sdnext/wiki/Video) for details! @@ -47,7 +47,7 @@ Plus support for **CogView-4**, **SANA 1.5**, new CLiP models, improvements to r - different video models support different video resolutions, frame counts, etc. and may require specific settings - see model links for details - see *ToDo/Limitations* section for additional notes -- **Models** +- **Models & Pipelines** - [THUDM CogView 4](https://huggingface.co/THUDM/CogView4-6B) **6B** variant new foundation model for image generation based o GLM-4 text encoder and a flow-based diffusion transformer fully supports offloading and on-the-fly quantization @@ -57,6 +57,11 @@ Plus support for **CogView-4**, **SANA 1.5**, new CLiP models, improvements to r big update to previous SANA model fully supports offloading and on-the-fly quantization simply select from *networks -> models -> reference* + - [ByteDance InfiniteYou](https://github.com/bytedance/InfiniteYou/): Flexible Photo Recrafting While Preserving Your Identity + face-transfer model for FLUX.1 + select from *Scripts -> InfiniteYou* + its large, ~12GB on top of FLUX.1 base model so make sure you have offloading and quantization setup + *note* model will be auto-downloaded on first use - New [zer0int CLiP-L](https://huggingface.co/zer0int/CLIP-Registers-Gated_MLP-ViT-L-14) models: download text encoders into folder set in settings -> system paths -> text encoders (default is *models/Text-encoder*) load using *settings -> text encoder* diff --git a/TODO.md b/TODO.md index 1ccc0a2dc..d9debcb5d 100644 --- a/TODO.md +++ b/TODO.md @@ -6,8 +6,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma ### Issues/Limitations -- VLM Gemma3: requires `transformers==git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3` -- VAE Remote encode: SD15 and Flux.1 issues: - Video: Hunyuan Video I2V: requires `transformers==4.47.1` - Video: Latte 1 T2V: dtype mismatch - Video: WAN 2.1 14B I2V 480p/720p: broken offload diff --git a/modules/infiniteyou/__init__.py b/modules/infiniteyou/__init__.py new file mode 100644 index 000000000..142921909 --- /dev/null +++ b/modules/infiniteyou/__init__.py @@ -0,0 +1,2 @@ +from .pipeline_flux_infusenet import FluxInfuseNetPipeline +from .pipeline_infu_flux import InfUFluxPipeline diff --git a/modules/infiniteyou/pipeline_flux_infusenet.py b/modules/infiniteyou/pipeline_flux_infusenet.py new file mode 100644 index 000000000..38fa186f9 --- /dev/null +++ b/modules/infiniteyou/pipeline_flux_infusenet.py @@ -0,0 +1,612 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. +# Copyright (c) 2024 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from diffusers import FluxControlNetPipeline +from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel +from diffusers.image_processor import PipelineImageInput +from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput +from diffusers.utils import is_torch_xla_available, logging + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.16, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class FluxInfuseNetPipeline(FluxControlNetPipeline): + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 28, + timesteps: List[int] = None, + guidance_scale: float = 3.5, + id_image: PipelineImageInput = None, + controlnet_guidance_scale: float = 1.0, + control_guidance_start: Union[float, List[float]] = 0.0, + control_guidance_end: Union[float, List[float]] = 1.0, + control_image: PipelineImageInput = None, + control_mode: Optional[Union[int, List[int]]] = None, + controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + + # ID-specific parameters + controlnet_prompt_embeds: Optional[torch.FloatTensor] = None, + + # True CFG parameters + true_guidance_scale: float = 1.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + will be used instead + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 7.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + controlnet_guidance_scale (`float`, *optional*, defaults to 7.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `controlnet_guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + The percentage of total steps at which the ControlNet starts applying. + control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + The percentage of total steps at which the ControlNet stops applying. + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: + `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + The ControlNet input condition to provide guidance to the `unet` for generation. If the type is + specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted + as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or + width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, + images must be passed as a list such that each element of the list can be correctly batched for input + to a single ControlNet. + controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added + to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set + the corresponding scale as a list. + control_mode (`int` or `List[int]`,, *optional*, defaults to None): + The control mode when applying ControlNet-Union. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + controlnet_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated embeddings for the InfuseNet. Can be used to easily tweak inputs, *e.g.* image embeddings. + If not provided, embeddings will be generated from `prompt` or `prompt_embeds` input arguments. + true_guidance_scale (`float`, *optional*, defaults to 1.0): + True CFG scale as defined in [Classifier-Free Diffusion Guidance]((https://arxiv.org/abs/2207.12598). + negative_prompt (`str` or `List[str]`, *optional*): + The negative prompt or negative prompts to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds`. instead. + negative_prompt_2 (`str` or `List[str]`, *optional*): + The negative prompt or negative prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, + `negative_prompt` is will be used instead. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative text embeddings will be generated from `negative_prompt` input + argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative pooled text embeddings will be generated from + `negative_prompt` input argument. + + Examples: + + Returns: + [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict` + is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated + images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): + control_guidance_start = len(control_guidance_end) * [control_guidance_start] + elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): + control_guidance_end = len(control_guidance_start) * [control_guidance_end] + elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): + mult = len(self.controlnet.nets) if isinstance(self.controlnet, FluxMultiControlNetModel) else 1 + control_guidance_start, control_guidance_end = ( + mult * [control_guidance_start], + mult * [control_guidance_end], + ) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + height, + width, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._controlnet_guidance_scale = controlnet_guidance_scale + self._true_guidance_scale = true_guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + dtype = self.transformer.dtype + + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) + ( + prompt_embeds, + pooled_prompt_embeds, + text_ids, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + if negative_prompt is not None or (negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None): + ( + negative_prompt_embeds, + negative_pooled_prompt_embeds, + negative_text_ids, + ) = self.encode_prompt( + prompt=negative_prompt, + prompt_2=negative_prompt_2, + prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + if controlnet_prompt_embeds is None: + controlnet_prompt_embeds = prompt_embeds + ( + controlnet_prompt_embeds, + pooled_prompt_embeds, + controlnet_text_ids, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_embeds=controlnet_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + # 3. Prepare control image + num_channels_latents = self.transformer.config.in_channels // 4 + if isinstance(self.controlnet, FluxControlNetModel) or True: + control_image = self.prepare_image( + image=control_image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=self.vae.dtype, + ) + height, width = control_image.shape[-2:] + + # xlab controlnet has a input_hint_block and instantx controlnet does not + controlnet_blocks_repeat = False if self.controlnet.input_hint_block is None else True + if self.controlnet.input_hint_block is None: + # vae encode + control_image = self.vae.encode(control_image).latent_dist.sample() + control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + # pack + height_control_image, width_control_image = control_image.shape[2:] + control_image = self._pack_latents( + control_image, + batch_size * num_images_per_prompt, + num_channels_latents, + height_control_image, + width_control_image, + ) + + # Here we ensure that `control_mode` has the same length as the control_image. + if control_mode is not None: + if not isinstance(control_mode, int): + raise ValueError(" For `FluxControlNet`, `control_mode` should be an `int` or `None`") + control_mode = torch.tensor(control_mode).to(device, dtype=torch.long) + control_mode = control_mode.view(-1, 1).expand(control_image.shape[0], 1) + + elif isinstance(self.controlnet, FluxMultiControlNetModel): + control_images = [] + # xlab controlnet has a input_hint_block and instantx controlnet does not + controlnet_blocks_repeat = False if self.controlnet.nets[0].input_hint_block is None else True + for _i, control_image_ in enumerate(control_image): + control_image_ = self.prepare_image( + image=control_image_, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=self.vae.dtype, + ) + height, width = control_image_.shape[-2:] + + if self.controlnet.nets[0].input_hint_block is None: + # vae encode + control_image_ = self.vae.encode(control_image_).latent_dist.sample() + control_image_ = (control_image_ - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + # pack + height_control_image, width_control_image = control_image_.shape[2:] + control_image_ = self._pack_latents( + control_image_, + batch_size * num_images_per_prompt, + num_channels_latents, + height_control_image, + width_control_image, + ) + control_images.append(control_image_) + + control_image = control_images + + # Here we ensure that `control_mode` has the same length as the control_image. + if isinstance(control_mode, list) and len(control_mode) != len(control_image): + raise ValueError("For Multi-ControlNet, `control_mode` must be a list of the same length as the number of controlnets (control images) specified") + if not isinstance(control_mode, list): + control_mode = [control_mode] * len(control_image) + # set control mode + control_modes = [] + for cmode in control_mode: + if cmode is None: + cmode = -1 + control_mode = torch.tensor(cmode).expand(control_images[0].shape[0]).to(device, dtype=torch.long) + control_modes.append(control_mode) + control_mode = control_modes + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels // 4 + latents, latent_image_ids = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + image_seq_len = latents.shape[1] + mu = calculate_shift( + image_seq_len, + self.scheduler.config.base_image_seq_len, + self.scheduler.config.max_image_seq_len, + self.scheduler.config.base_shift, + self.scheduler.config.max_shift, + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + mu=mu, + ) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Create tensor stating which controlnets to keep + controlnet_keep = [] + for i in range(len(timesteps)): + keeps = [ + 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) + for s, e in zip(control_guidance_start, control_guidance_end) + ] + controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps) + + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + if isinstance(self.controlnet, FluxMultiControlNetModel): + use_guidance = self.controlnet.nets[0].config.guidance_embeds + else: + use_guidance = self.controlnet.config.guidance_embeds + + guidance = torch.tensor([controlnet_guidance_scale], device=device) if use_guidance else None + guidance = guidance.expand(latents.shape[0]) if guidance is not None else None + + if isinstance(controlnet_keep[i], list): + if not isinstance(controlnet_conditioning_scale, list): + controlnet_conditioning_scale = len(controlnet_keep) * [controlnet_conditioning_scale] + cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] + controlnet_conditioning_scale = controlnet_conditioning_scale[0] + else: + controlnet_cond_scale = controlnet_conditioning_scale + if isinstance(controlnet_cond_scale, list): + controlnet_cond_scale = controlnet_cond_scale[0] + cond_scale = controlnet_cond_scale * controlnet_keep[i] + + # controlnet + controlnet_block_samples, controlnet_single_block_samples = self.controlnet( + hidden_states=latents, + controlnet_cond=control_image, + controlnet_mode=control_mode, + conditioning_scale=cond_scale[0], + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=controlnet_prompt_embeds, + txt_ids=controlnet_text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + ) + + guidance = ( + torch.tensor([guidance_scale], device=device) if self.transformer.config.guidance_embeds else None + ) + guidance = guidance.expand(latents.shape[0]) if guidance is not None else None + + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + controlnet_block_samples=controlnet_block_samples, + controlnet_single_block_samples=controlnet_single_block_samples, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + controlnet_blocks_repeat=controlnet_blocks_repeat, + )[0] + + # perform true CFG + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None and negative_text_ids is not None: + noise_pred_uncond = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + txt_ids=negative_text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + controlnet_blocks_repeat=controlnet_blocks_repeat, + )[0] + + noise_pred = noise_pred_uncond + true_guidance_scale * (noise_pred - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return FluxPipelineOutput(images=image) diff --git a/modules/infiniteyou/pipeline_infu_flux.py b/modules/infiniteyou/pipeline_infu_flux.py new file mode 100644 index 000000000..8ae9f6e95 --- /dev/null +++ b/modules/infiniteyou/pipeline_infu_flux.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os +import random +from typing import Optional + +import cv2 +import numpy as np +import torch +from diffusers.models import FluxControlNetModel +from facexlib.recognition import init_recognition_model +from huggingface_hub import snapshot_download +from insightface.app import FaceAnalysis +from insightface.utils import face_align +from PIL import Image + +from modules import shared, devices, model_quant +from .pipeline_flux_infusenet import FluxInfuseNetPipeline +from .resampler import Resampler + + +def seed_everything(seed, deterministic=False): + """Set random seed. + + Args: + seed (int): Seed to be used. + deterministic (bool): Whether to set the deterministic option for + CUDNN backend, i.e., set `torch.backends.cudnn.deterministic` + to True and `torch.backends.cudnn.benchmark` to False. + Default: False. + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + os.environ['PYTHONHASHSEED'] = str(seed) + if deterministic: + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def retrieve_latents( + encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# modified from https://github.com/instantX-research/InstantID/blob/main/pipeline_stable_diffusion_xl_instantid.py +def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]): + stickwidth = 4 + limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) + kps = np.array(kps) + + w, h = image_pil.size + out_img = np.zeros([h, w, 3]) + + for i in range(len(limbSeq)): + index = limbSeq[i] + color = color_list[index[0]] + + x = kps[index][:, 0] + y = kps[index][:, 1] + length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) + polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) + out_img = (out_img * 0.6).astype(np.uint8) + + for idx_kp, kp in enumerate(kps): + color = color_list[idx_kp] + x, y = kp + out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) + + out_img_pil = Image.fromarray(out_img.astype(np.uint8)) + return out_img_pil + + +def extract_arcface_bgr_embedding(in_image, landmark, arcface_model=None, in_settings=None): # pylint: disable=unused-argument + kps = landmark + arc_face_image = face_align.norm_crop(in_image, landmark=np.array(kps), image_size=112) + arc_face_image = torch.from_numpy(arc_face_image).unsqueeze(0).permute(0,3,1,2) / 255. + arc_face_image = 2 * arc_face_image - 1 + arc_face_image = arc_face_image.cuda().contiguous() + if arcface_model is None: + arcface_model = init_recognition_model('arcface', device=devices.device) + face_emb = arcface_model(arc_face_image)[0] # [512], normalized + return face_emb + + +def resize_and_pad_image(source_img, target_img_size): + # Get original and target sizes + source_img_size = source_img.size + target_width, target_height = target_img_size + + # Determine the new size based on the shorter side of target_img + if target_width <= target_height: + new_width = target_width + new_height = int(target_width * (source_img_size[1] / source_img_size[0])) + else: + new_height = target_height + new_width = int(target_height * (source_img_size[0] / source_img_size[1])) + + # Resize the source image using LANCZOS interpolation for high quality + resized_source_img = source_img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + # Compute padding to center resized image + pad_left = (target_width - new_width) // 2 + pad_top = (target_height - new_height) // 2 + + # Create a new image with white background + padded_img = Image.new("RGB", target_img_size, (255, 255, 255)) + padded_img.paste(resized_source_img, (pad_left, pad_top)) + + return padded_img + + +class InfUFluxPipeline: + def __init__( + self, + pipe, + image_proj_num_tokens=8, + infu_flux_version='v1.0', + model_version='aes_stage2', + ): + + self.infu_flux_version = infu_flux_version + self.model_version = model_version + + # Load pipeline + local_path = snapshot_download(repo_id='ByteDance/InfiniteYou', cache_dir=shared.opts.hfcache_dir) + infiniteyou_path = os.path.join(local_path, f'infu_flux_{infu_flux_version}', model_version) + infusenet_path = os.path.join(infiniteyou_path, 'InfuseNetModel') + quant_args = model_quant.create_config() + # quant_args = {} + + self.infusenet = FluxControlNetModel.from_pretrained( + infusenet_path, + torch_dtype=devices.dtype, + **quant_args, + ) + + self.pipe = FluxInfuseNetPipeline( + vae=pipe.vae, + text_encoder=pipe.text_encoder, + text_encoder_2=pipe.text_encoder_2, + tokenizer=pipe.tokenizer, + tokenizer_2=pipe.tokenizer_2, + transformer=pipe.transformer, + scheduler=pipe.scheduler, + controlnet=self.infusenet, + ) + + # Load image proj model + num_tokens = image_proj_num_tokens + image_emb_dim = 512 + image_proj_model = Resampler( + dim=1280, + depth=4, + dim_head=64, + heads=20, + num_queries=num_tokens, + embedding_dim=image_emb_dim, + output_dim=4096, + ff_mult=4, + ) + image_proj_model_path = os.path.join(infiniteyou_path, 'image_proj_model.bin') + ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu") + image_proj_model.load_state_dict(ipm_state_dict['image_proj']) + del ipm_state_dict + image_proj_model.to(device=devices.device, dtype=devices.dtype) + image_proj_model.eval() + + self.image_proj_model = image_proj_model + + # Load face encoder + insightface_root_path = os.path.join(local_path, 'supports', 'insightface') + self.app_640 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx) + self.app_640.prepare(ctx_id=0, det_size=(640, 640)) + self.app_320 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx) + self.app_320.prepare(ctx_id=0, det_size=(320, 320)) + self.app_160 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx) + self.app_160.prepare(ctx_id=0, det_size=(160, 160)) + self.arcface_model = init_recognition_model('arcface', device=devices.device) + + def load_loras(self, loras): + names, scales = [],[] + for lora_path, lora_name, lora_scale in loras: + if lora_path != "": + print(f"loading lora {lora_path}") + self.pipe.load_lora_weights(lora_path, adapter_name = lora_name) + names.append(lora_name) + scales.append(lora_scale) + + if len(names) > 0: + self.pipe.set_adapters(names, adapter_weights=scales) + + def _detect_face(self, id_image_cv2): + face_info = self.app_640.get(id_image_cv2) + if len(face_info) > 0: + return face_info + + face_info = self.app_320.get(id_image_cv2) + if len(face_info) > 0: + return face_info + + face_info = self.app_160.get(id_image_cv2) + return face_info + + def __call__( + self, + prompt: str, + id_image: Image.Image, # PIL.Image.Image (RGB) + negative_prompt = None, + control_image: Optional[Image.Image] = None, # PIL.Image.Image (RGB) or None + width = 1024, + height = 1024, + seed = 42, + guidance_scale = 3.5, + controlnet_guidance_scale = 1.0, + num_inference_steps = 30, + infusenet_conditioning_scale = 1.0, + infusenet_guidance_start = 0.0, + infusenet_guidance_end = 1.0, + output_type = 'pil', + generator = None, + *args, **kwargs # pylint: disable=unused-argument + ): + # Extract ID embeddings + id_image_cv2 = cv2.cvtColor(np.array(id_image), cv2.COLOR_RGB2BGR) + face_info = self._detect_face(id_image_cv2) + if len(face_info) == 0: + raise ValueError('No face detected in the input ID image') + + face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face + landmark = face_info['kps'] + id_embed = extract_arcface_bgr_embedding(id_image_cv2, landmark, self.arcface_model) + id_embed = id_embed.clone().unsqueeze(0).float().cuda() + id_embed = id_embed.reshape([1, -1, 512]) + id_embed = id_embed.to(device=devices.device, dtype=devices.dtype) + with torch.no_grad(): + id_embed = self.image_proj_model(id_embed) + bs_embed, seq_len, _ = id_embed.shape + id_embed = id_embed.repeat(1, 1, 1) + id_embed = id_embed.view(bs_embed * 1, seq_len, -1) + id_embed = id_embed.to(device=devices.device, dtype=devices.dtype) + + # Load control image + if control_image is not None: + control_image = control_image.convert("RGB") + control_image = resize_and_pad_image(control_image, (width, height)) + face_info = self._detect_face(cv2.cvtColor(np.array(control_image), cv2.COLOR_RGB2BGR)) + if len(face_info) == 0: + raise ValueError('No face detected in the control image') + face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face + control_image = draw_kps(control_image, face_info['kps']) + else: + out_img = np.zeros([height, width, 3]) + control_image = Image.fromarray(out_img.astype(np.uint8)) + + """ + control_image = self.pipe.prepare_image( + image=control_image, + width=width, + height=height, + batch_size=1, + num_images_per_prompt=1, + device=devices.device, + dtype=devices.dtype, + ) + control_image = retrieve_latents(self.pipe.vae.encode(control_image), generator=generator) + control_image = (control_image - self.pipe.vae.config.shift_factor) * self.pipe.vae.config.scaling_factor + # pack + height_control_image, width_control_image = control_image.shape[2:] + num_channels_latents = self.pipe.transformer.config.in_channels // 4 + control_image = self.pipe._pack_latents( + control_image, + 1, + num_channels_latents, + height_control_image, + width_control_image, + ) + """ + + # Perform inference + seed_everything(seed) + latents = self.pipe( + prompt=prompt, + negative_prompt=negative_prompt, + controlnet_prompt_embeds=id_embed, + control_image=control_image, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + controlnet_guidance_scale=controlnet_guidance_scale, + controlnet_conditioning_scale=infusenet_conditioning_scale, + control_guidance_start=infusenet_guidance_start, + control_guidance_end=infusenet_guidance_end, + height=height, + width=width, + output_type=output_type, + callback_on_step_end=kwargs.get('callback_on_step_end', None), + callback_on_step_end_tensor_inputs=kwargs.get('callback_on_step_end_tensor_inputs', None), + ) + + return latents diff --git a/modules/infiniteyou/resampler.py b/modules/infiniteyou/resampler.py new file mode 100644 index 000000000..6d0011e83 --- /dev/null +++ b/modules/infiniteyou/resampler.py @@ -0,0 +1,121 @@ +# Modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py + +import math + +import torch +import torch.nn as nn + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index b22825639..b1af5d9b6 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -485,7 +485,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn weights_backup = getattr(self, "network_weights_backup", None) bias_backup = getattr(self, "network_bias_backup", None) if weights_backup is None and bias_backup is None: - return None, None + return t0 = time.time() if weights_backup is not None: diff --git a/modules/processing_args.py b/modules/processing_args.py index db8794214..a839c9995 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -119,16 +119,16 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t t0 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) apply_circular(p.tiling, model) - if hasattr(model, "set_progress_bar_config"): - if disable_pbar: - model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba', disable=disable_pbar) - else: - model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba') args = {} has_vae = hasattr(model, 'vae') or (hasattr(model, 'pipe') and hasattr(model.pipe, 'vae')) if hasattr(model, 'pipe') and not hasattr(model, 'no_recurse'): # recurse model = model.pipe has_vae = has_vae or hasattr(model, 'vae') + if hasattr(model, "set_progress_bar_config"): + if disable_pbar: + model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba', disable=disable_pbar) + else: + model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba') signature = inspect.signature(type(model).__call__, follow_wrapped=True) possible = list(signature.parameters) diff --git a/scripts/infiniteyou_ext.py b/scripts/infiniteyou_ext.py new file mode 100644 index 000000000..59432cd40 --- /dev/null +++ b/scripts/infiniteyou_ext.py @@ -0,0 +1,121 @@ +# https://huggingface.co/ByteDance/InfiniteYou +# https://github.com/bytedance/InfiniteYou +# flux base model + 11.8gb controlnet module + 338mb image module + 428 insightface module + +import gradio as gr +from PIL import Image +from modules import scripts, processing, shared, sd_models, devices + + +prefix = 'InfiniteYou' +model_versions = ['aes_stage2', 'sim_stage1'] +orig_pipeline, orig_prompt_attention = None, None + + +def verify_insightface(): + from installer import installed, install, reload + if not installed('insightface', reload=False, quiet=True): + install('insightface==0.7.3', ignore=False) + install('albumentations==1.4.3', ignore=False, reinstall=True) + install('pydantic==1.10.21', ignore=False, reinstall=True, force=True) + reload('pydantic') + + +def load_infiniteyou(model: str): + from modules.infiniteyou import InfUFluxPipeline + shared.sd_model = InfUFluxPipeline( + pipe=shared.sd_model, + model_version=model, + ) + sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) + sd_models.set_diffuser_options(shared.sd_model) + + +class Script(scripts.Script): + def title(self): + return f'{prefix}: Flexible Photo Recrafting' + + def show(self, is_img2img): + return not is_img2img if shared.native else False + + # return signature is array of gradio components + def ui(self, _is_img2img): + with gr.Row(): + gr.HTML(f'  {prefix}: Flexible Photo Recrafting
') + with gr.Row(): + model = gr.Dropdown(label='IY model', choices=model_versions, value=model_versions[0]) + restore = gr.Checkbox(label='Restore pipeline on end', value=False) + with gr.Row(): + scale = gr.Slider(label='IY scale', value=1.0, minimum=0.0, maximum=2.0, step=0.05) + with gr.Row(): + start = gr.Slider(label='IY start', value=0.0, minimum=0.0, maximum=1.0, step=0.05) + end = gr.Slider(label='IY end', value=1.0, minimum=0.0, maximum=1.0, step=0.05) + with gr.Row(): + id_guidance = gr.Slider(label='Identity guidance', value=3.5, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + id_image = gr.Image(label='Identity image', type='pil') + with gr.Row(): + control_guidance = gr.Slider(label='Control guidance', value=1.0, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + control_image = gr.Image(label='Control image', type='pil') + return [model, id_image, control_image, scale, start, end, id_guidance, control_guidance, restore] + + def run(self, p: processing.StableDiffusionProcessing, + model: str = None, + id_image: Image.Image = None, + control_image: Image.Image = None, + scale: float = 1.0, + start: float = 0.0, + end: float = 1.0, + id_guidance: float = 3.5, + control_guidance: float = 1.0, + restore: bool = False, + ): # pylint: disable=arguments-differ, unused-argument + + if model is None or model not in model_versions: + return None + if id_image is None: + shared.log.error(f'{prefix}: no init_images') + return None + if shared.sd_model_type != 'f1': + shared.log.error(f'{prefix}: invalid model type: {shared.sd_model_type}') + return None + + global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement + orig_pipeline = shared.sd_model + if shared.sd_model.__class__.__name__ != 'InfUFluxPipeline': + verify_insightface() + load_infiniteyou(model) + devices.torch_gc() + shared.log.info(f'{prefix}: cls={shared.sd_model.__class__.__name__} loaded') + + processing.fix_seed(p) + p.task_args['id_image'] = id_image + p.task_args['control_image'] = control_image + p.task_args['infusenet_conditioning_scale'] = scale + p.task_args['infusenet_guidance_start'] = start + p.task_args['infusenet_guidance_end'] = end + p.task_args['seed'] = p.seed + p.task_args['negative_prompt'] = None + p.task_args['guidance_scale'] = id_guidance + p.task_args['controlnet_guidance_scale'] = control_guidance + p.extra_generation_params['IY model'] = model + p.extra_generation_params['IY guidance'] = f'{scale:.1f}/{start:.1f}/{end:.1f}' + orig_prompt_attention = shared.opts.prompt_attention + shared.opts.data['prompt_attention'] = 'fixed' + shared.log.debug(f'{prefix}: args={p.task_args}') + + processed = processing.process_images(p) + return processed + + def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args, **kwargs): # pylint: disable=unused-argument + # restore pipeline + global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement + restore = args[-1] + if orig_prompt_attention is not None: + shared.opts.data['prompt_attention'] = orig_prompt_attention + orig_prompt_attention = None + if restore and orig_pipeline is not None: + shared.log.info(f'{prefix}: restoring pipeline') + shared.sd_model = orig_pipeline + orig_pipeline = None