feat(anima): add image-to-image and inpainting support

modules/modeldata.py

@@ -72,7 +72,7 @@ def get_model_type(pipe):
         model_type = 'sana'
     elif "HiDream" in name:
         model_type = 'h1'
-    elif "AnimaTextToImage" in name:
+    elif name.startswith("Anima") and "AnimateDiff" not in name:
         model_type = 'anima'
     elif "Cosmos2TextToImage" in name:
         model_type = 'cosmos'

modules/processing_args.py

@@ -272,7 +272,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
             kwargs['output_type'] = 'np' # only set latent if model has vae
 
     # model specific
-    if 'Kandinsky' in model.__class__.__name__ or 'Cosmos2' in model.__class__.__name__ or 'Anima' in model.__class__.__name__ or 'OmniGen2' in model.__class__.__name__:
+    if 'Kandinsky' in model.__class__.__name__ or 'Cosmos2' in model.__class__.__name__ or 'OmniGen2' in model.__class__.__name__:
         kwargs['output_type'] = 'np' # only set latent if model has vae
     if 'StableCascade' in model.__class__.__name__:
         kwargs.pop("guidance_scale") # remove

modules/processing_vae.py

@@ -247,6 +247,8 @@ def vae_postprocess(tensor, model, output_type='np'):
                 if tensor.ndim == 6 and tensor.shape[1] == 1:
                     tensor = tensor.squeeze(0)
                 images = model.video_processor.postprocess_video(tensor, output_type='pil')
+                if isinstance(images, list) and len(images) > 0 and isinstance(images[0], list):
+                    images = [frame for batch in images for frame in batch]
             elif hasattr(model, 'image_processor'):
                 if tensor.ndim == 5 and tensor.shape[1] == 3: # Qwen Image
                     tensor = tensor[:, :, 0]

pipelines/anima/anima_image.py

@@ -0,0 +1,168 @@
+"""Anima img2img and inpainting pipelines (built dynamically from the runtime-imported base class)."""
+
+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
+
+
+def _encode_image(pipe, image, dtype, device, generator):
+    """VAE-encode an image and normalize to denoiser latent space."""
+    if isinstance(image, list):
+        image = image[0]
+    image_tensor = pipe.video_processor.preprocess(image, None, None)
+    image_tensor = image_tensor.squeeze(0).to(device=device, dtype=pipe.vae.dtype)
+    image_tensor = image_tensor[None, :, None, :, :]
+    init_latents = pipe.vae.encode(image_tensor).latent_dist.sample(generator)
+    latents_mean = torch.tensor(pipe.vae.config.latents_mean, device=device, dtype=torch.float32).view(1, pipe.vae.config.z_dim, 1, 1, 1)
+    latents_std_inv = (1.0 / torch.tensor(pipe.vae.config.latents_std, device=device, dtype=torch.float32)).view(1, pipe.vae.config.z_dim, 1, 1, 1)
+    return ((init_latents.float() - latents_mean) * latents_std_inv).to(dtype)
+
+
+def _setup_img2img_schedule(scheduler, strength, num_inference_steps, device):
+    """Set custom sigma schedule, return first sigma after scheduler shift."""
+    custom_sigmas = torch.linspace(max(strength, 0.01), 0.0, num_inference_steps).tolist()
+    scheduler.set_timesteps(sigmas=custom_sigmas, device=device)
+    return scheduler.sigmas[0].item()
+
+
+def build_anima_pipeline_classes(base_cls):
+    """Return (AnimaImageToImagePipeline, AnimaInpaintPipeline) inheriting from base_cls."""
+
+    class AnimaImageToImagePipeline(base_cls):
+        """Anima img2img 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,
+            strength: float = 0.8,
+            height: int = 768,
+            width: int = 1360,
+            num_inference_steps: int = 35,
+            guidance_scale: float = 7.0,
+            num_images_per_prompt: Optional[int] = 1,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.Tensor] = None,
+            prompt_embeds: Optional[torch.Tensor] = None,
+            negative_prompt_embeds: Optional[torch.Tensor] = None,
+            output_type: Optional[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"],
+            max_sequence_length: int = 512,
+        ):
+            actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+            init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+            noise = randn_tensor(init_latents.shape, generator=generator, device=devices.device, dtype=devices.dtype)
+            noised = (actual_sigma * noise + (1.0 - actual_sigma) * init_latents).to(torch.float32)
+
+            orig_set_timesteps = self.scheduler.set_timesteps
+            self.scheduler.set_timesteps = lambda *args, **kwargs: None
+            try:
+                return super().__call__(
+                    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=callback_on_step_end, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+                    max_sequence_length=max_sequence_length,
+                )
+            finally:
+                self.scheduler.set_timesteps = orig_set_timesteps
+
+    class AnimaInpaintPipeline(AnimaImageToImagePipeline):
+        """Anima 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 = 768,
+            width: int = 1360,
+            num_inference_steps: int = 35,
+            guidance_scale: float = 7.0,
+            num_images_per_prompt: Optional[int] = 1,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.Tensor] = None,
+            prompt_embeds: Optional[torch.Tensor] = None,
+            negative_prompt_embeds: Optional[torch.Tensor] = None,
+            output_type: Optional[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"],
+            max_sequence_length: int = 512,
+        ):
+            actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, devices.device)
+            init_latents = _encode_image(self, image, devices.dtype, devices.device, generator)
+            noise = randn_tensor(init_latents.shape, generator=generator, device=devices.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, devices.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 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,
+                    callback_on_step_end_tensor_inputs=["latents"],
+                    max_sequence_length=max_sequence_length,
+                )
+            finally:
+                self.scheduler.set_timesteps = orig_set_timesteps
+
+    return AnimaImageToImagePipeline, AnimaInpaintPipeline
+
+
+def _prepare_mask(pipe, mask_image, height, width, device):
+    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_scale_factor_spatial
+    latent_w = width // pipe.vae_scale_factor_spatial
+    mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest")
+    mask_latent = mask_latent[:, :1, :, :]
+    mask_latent = mask_latent.unsqueeze(2)
+    return mask_latent

pipelines/model_anima.py

@@ -81,6 +81,12 @@ def load_anima(checkpoint_info, diffusers_load_config=None):
     AnimaTextToImagePipeline = pipeline_mod.AnimaTextToImagePipeline
     AnimaLLMAdapter = adapter_mod.AnimaLLMAdapter
 
+    from pipelines.anima.anima_image import build_anima_pipeline_classes
+    AnimaImageToImagePipeline, AnimaInpaintPipeline = build_anima_pipeline_classes(AnimaTextToImagePipeline)
+    diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["anima"] = AnimaTextToImagePipeline
+    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["anima"] = AnimaImageToImagePipeline
+    diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["anima"] = AnimaInpaintPipeline
+
     # UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a
     # community file that bundles both the transformer and the llm_adapter.
     transformer, llm_adapter = load_transformer_components(repo_id, diffusers_load_config, AnimaLLMAdapter)