diff --git a/CHANGELOG.md b/CHANGELOG.md index f74446a6d..14c0d6cc3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,9 +13,11 @@ - [Nunchaku-Lite](https://huggingface.co/lite-infer) pre-quantized models included: *Z-Image, Flux.1-Dev/Schnell/Krea/Kontex, Qwen-Image/Image-Edit, Ernie-Image* - **Features** - - add support for [Nunchaku-Lite](https://github.com/rootonchair/nunchaku-lite) inference engine + - [Nunchaku-Lite](https://github.com/rootonchair/nunchaku-lite) inference engine unlike Nunchaku, Nunchaku-Lite is based on Kernels and does not require any additional packages to be installed but like original Nunchaku, it is only available for CUDA and right now only for `torch==2.11/2.12` + - [Krea2](https://huggingface.co/krea/Krea-2-Turbo) add *Inpaint* pipeline + this also makes K2 compatible with *Detailer* workflow - storage analyzer: new feature that analyzes your storage used by sdnext per type and location *system -> storage* - video: support for scripts/extensions diff --git a/pipelines/anima/anima_image.py b/pipelines/anima/anima_image.py index 5d0334181..233a35a96 100644 --- a/pipelines/anima/anima_image.py +++ b/pipelines/anima/anima_image.py @@ -131,12 +131,20 @@ def build_anima_pipeline_classes(base_cls): try: return base_cls.__call__( self, - prompt=prompt, negative_prompt=negative_prompt, - height=height, width=width, num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, num_images_per_prompt=num_images_per_prompt, - generator=generator, latents=noised, prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, output_type=output_type, - return_dict=return_dict, callback_on_step_end=blend_callback, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + num_images_per_prompt=num_images_per_prompt, + generator=generator, + latents=noised, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + output_type=output_type, + return_dict=return_dict, + callback_on_step_end=blend_callback, callback_on_step_end_tensor_inputs=["latents"], max_sequence_length=max_sequence_length, ) diff --git a/pipelines/krea2/pipeline_krea2_inpaint.py b/pipelines/krea2/pipeline_krea2_inpaint.py new file mode 100644 index 000000000..30533a285 --- /dev/null +++ b/pipelines/krea2/pipeline_krea2_inpaint.py @@ -0,0 +1,132 @@ +"""Krea 2 inpainting pipeline. + +This module adds inpainting support on top of the existing Krea2 image-to-image +variant without modifying the upstream Krea2 denoising path. +""" + +from typing import Callable, Dict, List, Optional, Union + +import torch +import torch.nn.functional as F +from PIL import Image +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import PipelineImageInput +from diffusers.utils.torch_utils import randn_tensor + +from modules import devices +from .pipeline_krea2 import Krea2Pipeline, Krea2Img2ImgPipeline + + +def _setup_img2img_schedule(scheduler, strength, num_inference_steps, device, mu=None): + """Set custom sigma schedule, return first sigma after scheduler shift.""" + min_sigma = 1e-8 + custom_sigmas = torch.linspace(max(strength, 0.01), min_sigma, num_inference_steps).tolist() + scheduler.set_timesteps(sigmas=custom_sigmas, device=device, mu=mu) + return scheduler.sigmas[0].item() + + +def _prepare_mask(pipe, mask_image, height, width, device): + if isinstance(mask_image, list): + mask_image = mask_image[0] + if isinstance(mask_image, Image.Image): + mask_image = mask_image.convert("L") + if isinstance(mask_image, Image.Image): + import torchvision.transforms.functional as TF + + mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).to(device=device, dtype=torch.float32) + elif isinstance(mask_image, torch.Tensor): + mask_tensor = mask_image.to(device=device, dtype=torch.float32) + if mask_tensor.ndim == 2: + mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0) + elif mask_tensor.ndim == 3: + mask_tensor = mask_tensor.unsqueeze(0) + else: + mask_tensor = torch.ones(1, 1, height, width, device=device, dtype=torch.float32) + + latent_h = height // pipe.vae_compression + latent_w = width // pipe.vae_compression + mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest") + return mask_latent[:, :1, :, :] + + +class Krea2InpaintPipeline(Krea2Img2ImgPipeline): + """Krea 2 inpainting pipeline.""" + + @torch.no_grad() + def __call__( + self, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + image: Optional[PipelineImageInput] = None, + mask_image: Optional[PipelineImageInput] = None, + strength: float = 0.8, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 28, + guidance_scale: float | None = None, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + ): + align = self.vae_compression * self.patch + height = (height // align) * align + width = (width // align) * align + + device = devices.device + dtype = self.transformer.dtype + + cfg = self.scheduler.config + grid_h = height // (self.vae_compression * self.patch) + grid_w = width // (self.vae_compression * self.patch) + mu = self.calculate_shift( + grid_h * grid_w, + cfg.get("base_image_seq_len", 256), + cfg.get("max_image_seq_len", 6400), + cfg.get("base_shift", 0.5), + cfg.get("max_shift", 1.15), + ) + actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, device, mu=mu) + init_latents = self.encode_image(image, height, width, dtype, device) + noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=devices.dtype) + noised = (actual_sigma * noise + (1.0 - actual_sigma) * init_latents).to(torch.float32) + mask_latent = _prepare_mask(self, mask_image, height, width, device) + + orig_set_timesteps = self.scheduler.set_timesteps + self.scheduler.set_timesteps = lambda *args, **kwargs: None + + user_callback = callback_on_step_end + + def blend_callback(pipe, i, t, callback_kwargs): + cur_latents = callback_kwargs.get("latents") + if cur_latents is not None: + sigma_next = pipe.scheduler.sigmas[i + 1].item() if i + 1 < len(pipe.scheduler.sigmas) else 0.0 + init_at_t = sigma_next * noise + (1.0 - sigma_next) * init_latents + blended = mask_latent * cur_latents + (1.0 - mask_latent) * init_at_t.to(cur_latents.dtype) + callback_kwargs["latents"] = blended + if user_callback is not None: + callback_kwargs = user_callback(pipe, i, t, callback_kwargs) + return callback_kwargs + + try: + return Krea2Pipeline.__call__( + self, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + num_images_per_prompt=num_images_per_prompt, + generator=generator, + latents=noised, + output_type=output_type, + return_dict=return_dict, + callback_on_step_end=blend_callback, + callback_on_step_end_tensor_inputs=["latents"], + ) + finally: + self.scheduler.set_timesteps = orig_set_timesteps diff --git a/pipelines/model_krea2.py b/pipelines/model_krea2.py index afe3ec5ac..4c65fb2e1 100644 --- a/pipelines/model_krea2.py +++ b/pipelines/model_krea2.py @@ -15,16 +15,19 @@ def load_krea2(checkpoint_info, diffusers_load_config=None): from pipelines.krea2.transformer_krea2 import Krea2Transformer2DModel from pipelines.krea2.pipeline_krea2 import Krea2Pipeline, Krea2Img2ImgPipeline + from pipelines.krea2.pipeline_krea2_inpaint import Krea2InpaintPipeline from pipelines.krea2 import KREA2_SPEC diffusers.Krea2Transformer2DModel = Krea2Transformer2DModel diffusers.Krea2Pipeline = Krea2Pipeline diffusers.Krea2Img2ImgPipeline = Krea2Img2ImgPipeline + diffusers.Krea2InpaintPipeline = Krea2InpaintPipeline generic.set_pipeline('Krea2', Krea2Pipeline) # One class per task so get_diffusers_task defaults to text2image and set_diffuser_pipe switches # to the img2img variant cleanly (matches the Chroma/Qwen per-task-class pattern). from diffusers.pipelines import auto_pipeline auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Pipeline auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Img2ImgPipeline + auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING['krea2'] = Krea2InpaintPipeline if repo_id is None or repo_id.lower() == 'none': return None