libcom 0.1.8__py3-none-any.whl

libcom/os_insert/source/sam_on_objectstitch.py

@@ -0,0 +1,253 @@
+"""SAM wrapper for ObjectStitch composites (skeleton).
+
+This module provides a libcom-style home for running SAM on top of
+ObjectStitch outputs, using a bbox prompt to focus on the insertion region.
+
+The actual SAM model loading (segment-anything) will be implemented in a
+later phase. For now only the public API is defined so that `OSInsertModel`
+can depend on it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SAM_PRETRAINED_ROOT = REPO_ROOT / "pretrained_models" / "sam"
+
+
+@dataclass
+class SamOnObjectStitchConfig:
+    """Configuration for SAM-on-ObjectStitch inference."""
+
+    sam_checkpoint: Path | str = SAM_PRETRAINED_ROOT / "sam_vit_h_4b8939.pth"
+    model_type: str = "vit_h"
+    device: str = "cuda:0"
+
+
+def _run_sam_on_image(predictor, image_bgr: np.ndarray, box_xyxy: Tuple[int, int, int, int]) -> np.ndarray:
+    """Run SAM with a box prompt and return a binary mask (uint8 0/255).
+
+    Ported from `run_sam_on_objectstitch.py::run_sam_on_image`.
+    """
+
+    image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    predictor.set_image(image)
+
+    box = np.array(box_xyxy, dtype=np.float32)
+    masks, scores, _ = predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=box[None, :],
+        multimask_output=True,
+    )
+    if masks is None or len(masks) == 0:
+        raise RuntimeError("SAM returned no masks")
+
+    best_idx = int(np.argmax(scores))
+    m = masks[best_idx].astype(np.uint8)
+
+    x1, y1, x2, y2 = box_xyxy
+    box_w = max(int(x2 - x1), 1)
+    box_h = max(int(y2 - y1), 1)
+    box_area = float(box_w * box_h)
+    area = float(int(m.sum()))
+    if area / box_area < 0.01:
+        raise RuntimeError("SAM mask area too small")
+
+    m = m.astype(np.uint8) * 255
+    return m
+
+
+def run_sam_on_objectstitch_single(
+    os_image_path: Path,
+    bg_shape_hw: Tuple[int, int],
+    bbox_xyxy_bg: Tuple[int, int, int, int],
+    *,
+    config: SamOnObjectStitchConfig,
+    out_path: Path | None = None,
+) -> Path:
+    """Run SAM on a single ObjectStitch composite and return a binary mask path.
+
+    Parameters
+    ----------
+    os_image_path:
+        Path to the ObjectStitch composite image (BGR when read by OpenCV).
+    bg_shape_hw:
+        Original background image shape as (H, W). The bbox is defined in this
+        coordinate system.
+    bbox_xyxy_bg:
+        Bbox (x1, y1, x2, y2) in the original background coordinates.
+    config:
+        SAM configuration including checkpoint and device.
+    out_path:
+        Optional output path for the SAM mask PNG. If None, will create a
+        sibling file next to the ObjectStitch input with suffix `_sam_mask.png`.
+    """
+
+    os_image_path = Path(os_image_path)
+    if out_path is None:
+        out_path = os_image_path.with_name(os_image_path.stem + "_sam_mask.png")
+    out_path = Path(out_path)
+
+    if not os_image_path.exists():
+        raise FileNotFoundError(os_image_path)
+    if not config.sam_checkpoint.exists():
+        raise FileNotFoundError(
+            f"SAM checkpoint not found: {config.sam_checkpoint}. "
+            f"Expected at {SAM_PRETRAINED_ROOT}."
+        )
+
+    os_image_bgr = cv2.imread(str(os_image_path))
+    if os_image_bgr is None:
+        raise RuntimeError(f"Failed to read ObjectStitch image: {os_image_path}")
+
+    os_h, os_w = os_image_bgr.shape[:2]
+    bg_h, bg_w = bg_shape_hw
+
+    bx1, by1, bx2, by2 = bbox_xyxy_bg
+
+    scale_x = os_w / float(bg_w) if bg_w > 0 else 1.0
+    scale_y = os_h / float(bg_h) if bg_h > 0 else 1.0
+
+    x1 = int(bx1 * scale_x)
+    y1 = int(by1 * scale_y)
+    x2 = int(bx2 * scale_x)
+    y2 = int(by2 * scale_y)
+
+    x1 = max(0, min(x1, os_w - 1))
+    y1 = max(0, min(y1, os_h - 1))
+    x2 = max(x1 + 1, min(x2, os_w))
+    y2 = max(y1 + 1, min(y2, os_h))
+
+    try:
+        from segment_anything import SamPredictor, sam_model_registry
+    except ImportError as e:
+        raise ImportError(
+            "segment_anything package is required for SAM aggressive mode. "
+            "Install from https://github.com/facebookresearch/segment-anything"
+        ) from e
+
+    sam = sam_model_registry[config.model_type](checkpoint=str(config.sam_checkpoint))
+    sam.to(device=config.device)
+    predictor = SamPredictor(sam)
+
+    os_mask = _run_sam_on_image(predictor, os_image_bgr, (x1, y1, x2, y2))
+    cv2.imwrite(str(out_path), os_mask)
+
+    return out_path
+
+
+def run_sam_on_objectstitch(
+    os_image: np.ndarray,
+    bg_shape_hw: Tuple[int, int],
+    bbox_xyxy_bg: Tuple[int, int, int, int],
+    *,
+    config: SamOnObjectStitchConfig,
+) -> np.ndarray:
+    os_image_bgr = os_image
+    if os_image_bgr.ndim != 3 or os_image_bgr.shape[2] != 3:
+        raise ValueError("Expected HxWx3 image")
+
+    os_h, os_w = os_image_bgr.shape[:2]
+    bg_h, bg_w = bg_shape_hw
+
+    bx1, by1, bx2, by2 = bbox_xyxy_bg
+
+    scale_x = os_w / float(bg_w) if bg_w > 0 else 1.0
+    scale_y = os_h / float(bg_h) if bg_h > 0 else 1.0
+
+    x1 = int(bx1 * scale_x)
+    y1 = int(by1 * scale_y)
+    x2 = int(bx2 * scale_x)
+    y2 = int(by2 * scale_y)
+
+    x1 = max(0, min(x1, os_w - 1))
+    y1 = max(0, min(y1, os_h - 1))
+    x2 = max(x1 + 1, min(x2, os_w))
+    y2 = max(y1 + 1, min(y2, os_h))
+
+    if not config.sam_checkpoint.exists():
+        raise FileNotFoundError(
+            f"SAM checkpoint not found: {config.sam_checkpoint}. "
+            f"Expected at {SAM_PRETRAINED_ROOT}."
+        )
+
+    try:
+        from segment_anything import SamPredictor, sam_model_registry
+    except ImportError as e:
+        raise ImportError(
+            "segment_anything package is required for SAM aggressive mode. "
+            "Install from https://github.com/facebookresearch/segment-anything"
+        ) from e
+
+    sam = sam_model_registry[config.model_type](checkpoint=str(config.sam_checkpoint))
+    sam.to(device=config.device)
+    predictor = SamPredictor(sam)
+
+    return _run_sam_on_image(predictor, os_image_bgr, (x1, y1, x2, y2))
+
+
+def build_sam_predictor(*, config: SamOnObjectStitchConfig):
+    # 检查 sam_checkpoint 是本地路径还是 HuggingFace 模型 ID
+    if isinstance(config.sam_checkpoint, Path):
+        if not config.sam_checkpoint.exists():
+            raise FileNotFoundError(
+                f"SAM checkpoint not found: {config.sam_checkpoint}. "
+                f"Expected at {SAM_PRETRAINED_ROOT}."
+            )
+        sam_checkpoint = str(config.sam_checkpoint)
+    else:
+        # 如果是字符串，可能是 HuggingFace 模型 ID，直接使用
+        sam_checkpoint = config.sam_checkpoint
+
+    try:
+        from segment_anything import SamPredictor, sam_model_registry
+    except ImportError as e:
+        raise ImportError(
+            "segment_anything package is required for SAM aggressive mode. "
+            "Install from https://github.com/facebookresearch/segment-anything"
+        ) from e
+    
+    sam = sam_model_registry[config.model_type](checkpoint=sam_checkpoint)
+    sam.to(device=config.device)
+    return SamPredictor(sam)
+
+
+def run_sam_on_objectstitch_with_predictor(
+    *,
+    predictor,
+    os_image: np.ndarray,
+    bg_shape_hw: Tuple[int, int],
+    bbox_xyxy_bg: Tuple[int, int, int, int],
+) -> np.ndarray:
+    os_image_bgr = os_image
+    if os_image_bgr.ndim != 3 or os_image_bgr.shape[2] != 3:
+        raise ValueError("Expected HxWx3 image")
+
+    os_h, os_w = os_image_bgr.shape[:2]
+    bg_h, bg_w = bg_shape_hw
+
+    bx1, by1, bx2, by2 = bbox_xyxy_bg
+
+    scale_x = os_w / float(bg_w) if bg_w > 0 else 1.0
+    scale_y = os_h / float(bg_h) if bg_h > 0 else 1.0
+
+    x1 = int(bx1 * scale_x)
+    y1 = int(by1 * scale_y)
+    x2 = int(bx2 * scale_x)
+    y2 = int(by2 * scale_y)
+
+    x1 = max(0, min(x1, os_w - 1))
+    y1 = max(0, min(y1, os_h - 1))
+    x2 = max(x1 + 1, min(x2, os_w))
+    y2 = max(y1 + 1, min(y2, os_h))
+
+    return _run_sam_on_image(predictor, os_image_bgr, (x1, y1, x2, y2))

libcom/os_insert/source/utils.py

@@ -0,0 +1,41 @@
+"""Shared utility functions for OSInsert."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+
+
+def load_bbox_txt(path: str | Path) -> Tuple[int, int, int, int]:
+    """Load a simple `x1 y1 x2 y2` bbox txt file."""
+    path = Path(path)
+    with path.open("r") as f:
+        line = f.readline().strip()
+    x1, y1, x2, y2 = map(int, line.split())
+    return x1, y1, x2, y2
+
+
+def make_rect_mask_from_bbox(h: int, w: int, bbox: Tuple[int, int, int, int]) -> np.ndarray:
+    """Create a uint8 mask (0/255) from bbox coordinates."""
+    x1, y1, x2, y2 = bbox
+    x1 = max(0, min(w - 1, x1))
+    x2 = max(0, min(w, x2))
+    y1 = max(0, min(h - 1, y1))
+    y2 = max(0, min(h, y2))
+    if x2 <= x1 or y2 <= y1:
+        raise ValueError("degenerate bbox")
+
+    mask = np.zeros((h, w), dtype=np.uint8)
+    mask[y1:y2, x1:x2] = 255
+    return mask
+
+
+def read_image_bgr(path: str | Path) -> np.ndarray:
+    path = Path(path)
+    img = cv2.imread(str(path))
+    if img is None:
+        raise FileNotFoundError(path)
+    return img

libcom/os_insert/source/v1.yaml

@@ -0,0 +1,108 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: libcom.os_insert.source.ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "composition"
+    cond_stage_key: "image"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: False   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+    u_cond_percent: 0.2
+    
+    scheduler_config: # 10000 warmup steps
+      target: libcom.os_insert.source.ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 1000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: libcom.os_insert.source.ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 9
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+        add_conv_in_front_of_unet: False
+
+    first_stage_config:
+      target: libcom.os_insert.source.ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 512
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: libcom.os_insert.source.ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
+      params:
+        version: openai-clip-vit-large-patch14
+
+data:
+  target: libcom.os_insert.source.ldm.data.datamodule.DataModuleFromConfig
+  dataset_dir: &dataset_dir path-to-openimages
+  params:
+      batch_size: 16
+      num_workers: 16
+      wrap: False
+      train:
+          target: libcom.os_insert.source.ldm.data.open_images.OpenImageDataset
+          params:
+              split: train
+              dataset_dir: *dataset_dir
+              image_size: 512
+      validation:
+          target: libcom.os_insert.source.ldm.data.open_images.COCOEEDataset
+          params:
+              split: test
+              dataset_dir: path-to-cocoee
+              image_size: 512
+      test:
+          target: libcom.os_insert.source.ldm.data.open_images.OpenImageDataset
+          params:
+              split: validation
+              dataset_dir: *dataset_dir
+              image_size: 512
+
+lightning:
+  trainer:
+    max_epochs: 40
+    accelerator: gpu
+    devices: 8
+    num_nodes: 1
+    val_check_interval: 1.0

libcom/painterly_image_harmonization/__init__.py

@@ -0,0 +1,4 @@
+# change to your lib name
+from .painterly_image_harmonization import PainterlyHarmonizationModel
+
+__all__ = ['PainterlyHarmonizationModel']

libcom/painterly_image_harmonization/painterly_image_harmonization.py

@@ -0,0 +1,215 @@
+import torch
+import torchvision
+from libcom.utils.model_download import download_pretrained_model, download_entire_folder
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+import torch 
+import torch.nn.functional as F
+import os
+import torchvision.transforms as transforms
+from tqdm.auto import tqdm
+from pytorch_lightning import seed_everything
+from omegaconf import OmegaConf
+from PIL import Image
+import torchvision.transforms.functional as tf
+from .source.PHDNet.phdnet import PHDNet
+from .source.PHDiffusion.ldm.modules.encoders.adapter import Adapter,NoRes_Adapter
+from .source.PHDiffusion.ldm.models.diffusion.scheduling_pndm import PNDMScheduler
+from .source.PHDiffusion.ldm.util import instantiate_from_config
+import logging
+logging.getLogger('transformers').setLevel(logging.ERROR) # disable transformer lib warning 
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['PHDNet', 'PHDiffusion'] 
+
+class PainterlyHarmonizationModel:
+    """
+    Painterly image harmonization prediction model.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type
+        kwargs (dict): use_residual (bool): whether to use adapter with residual or not for PHDiffusion
+
+    Examples:
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> from libcom import PainterlyHarmonizationModel
+        >>> import cv2
+        >>> import torch
+        >>> task_name = 'painterly_image_harmonization'
+        >>> MODEL_TYPE = 'PHDNet' # choose from 'PHDNet', 'PHDiffusion'
+        >>> comp_img = '../tests/painterly_harmonization_source/composite/3.png'
+        >>> comp_mask = '../tests/painterly_harmonization_source/composite_mask/3.png'
+        >>> net = PainterlyHarmonizationModel(device=0, model_type=MODEL_TYPE)
+        >>> output_img = net(comp_img, comp_mask)
+        >>> grid_img = make_image_grid([comp_img, comp_mask, output_img])
+        >>> cv2.imshow('painterly_image_harmonization_demo', grid_img)
+
+    Expected result:
+
+    .. image:: _static/image/painterly_image_harmonization_result2.jpg
+    .. image:: _static/image/painterly_image_harmonization_result3.jpg
+        
+    """
+    def __init__(self, device=0, model_type='PHDNet', **kwargs):
+        assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        
+        if model_type == 'PHDNet':
+            weight_path = os.path.join(model_dir, 'pretrained_models', model_type + '.pth')
+            self.device = check_gpu_device(device)
+            download_pretrained_model(weight_path)
+            self.build_pretrained_model(weight_path)
+        elif model_type == 'PHDiffusion':
+            self.use_residual = self.option.get('use_residual', True)
+            if self.use_residual:
+                phdiff_weight_path = os.path.join(model_dir, 'pretrained_models', model_type+'WithRes.pth')
+            else:
+                phdiff_weight_path = os.path.join(model_dir, 'pretrained_models', model_type+'.pth')
+            sd_weight_path = os.path.join(model_dir, '../shared_pretrained_models', 'sd-v1-4.ckpt')
+            download_pretrained_model(sd_weight_path)
+            download_pretrained_model(phdiff_weight_path)
+            self.device = check_gpu_device(device)
+            self.build_pretrained_model(sd_weight_path, phdiff_weight_path)
+        self.build_data_transformer()
+
+    def build_pretrained_model(self, *weight_path):
+        if len(weight_path) == 1:
+            weight_path = weight_path[0]
+            # build PHDNet
+            assert self.model_type == 'PHDNet', self.model_type
+            model = PHDNet(self.device)
+            model.load_networks(weight_path)
+            self.model = model.to(self.device).eval()
+        elif len(weight_path) == 2:
+            # build PHDiffusion model
+            sd_weight_path, phdiff_weight_path = weight_path
+            assert self.model_type == 'PHDiffusion', self.model_type
+            self.config= OmegaConf.load(cur_dir+'/source/PHDiffusion/stable_diffusion.yaml')
+            self.config.model.params.cond_stage_config.params.model_path = "openai/clip-vit-large-patch14"
+            pl_sd = torch.load(sd_weight_path, map_location="cpu", weights_only=False)
+            sd = pl_sd["state_dict"]
+            model = instantiate_from_config(self.config.model)
+            model.load_state_dict(sd, strict=False)
+
+            if self.use_residual:
+                adapter=Adapter(cin=int(64 * 4), channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False)
+            else:
+                adapter=NoRes_Adapter(cin=int(64 * 4), channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False)
+            model_resume_state = torch.load(phdiff_weight_path, map_location='cpu')
+            adapter.load_state_dict(model_resume_state['ad'])
+            model.model.diffusion_model.interact_blocks.load_state_dict(model_resume_state['interact'])
+
+            self.model   = model.to(self.device).eval()
+            self.adapter = adapter.to(self.device).eval()
+            
+            self.scheduler = PNDMScheduler(
+                                    beta_end=0.012,
+                                    beta_schedule='scaled_linear',
+                                    beta_start=0.00085,
+                                    num_train_timesteps=1000,
+                                    set_alpha_to_one=False,
+                                    skip_prk_steps=True,
+                                    steps_offset=1,
+                                    )
+
+    def build_data_transformer(self):
+        self.image_size = 512
+        self.transformer = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+        self.mask_transform = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+        ])
+    
+    def inputs_preprocess(self, composite_image, composite_mask):
+        img  = read_image_pil(composite_image)
+        img  = self.transformer(img)
+        mask = read_mask_pil(composite_mask).convert('L')
+        mask = self.mask_transform(mask)
+        if self.model_type == 'PHDiffusion':
+            img  = img.unsqueeze(0).to(self.device)
+            mask = mask.unsqueeze(0).to(self.device) 
+        return img, mask
+    
+    def outputs_postprocess(self, outputs):
+        if outputs.dim() == 4:
+            outputs = outputs.squeeze(0)
+        outputs = (torch.clamp((outputs.permute(1, 2, 0) + 1) / 2.0 * 255, 0, 255)).cpu().numpy()
+        outputs = outputs.astype(np.uint8)
+        outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
+        return outputs
+    
+    @torch.no_grad()
+    def PHDNet_inference(self, composite_image, composite_mask):
+        comp, mask = self.inputs_preprocess(composite_image, composite_mask)
+        outputs    = self.model(comp, mask)
+        preds      = self.outputs_postprocess(outputs)
+        return preds
+    
+    @torch.no_grad()
+    def __call__(self, composite_image, composite_mask, sample_steps=50, strength=0.7, random_seed=None):
+        """
+        Generating the harmonized image for the given composite image and the corresponding composite mask.
+
+        Args:
+            composite_image (str | numpy.ndarray): The path to the composite image or the composite image in ndarray form.
+            composite_mask (str | numpy.ndarray): The path to the composite mask or the composite mask in ndarray form.
+            sample_steps (int): Default total step in the inference process of PHDiffusion.
+            strength (float): A hyper-parameter that decides the total step (strength * sample_steps) for PHDiffusion.
+
+        Returns:
+            preds (numpy.ndarray): Generated harmonized image for the given composite image and the corresponding composite mask, with BGR channel.
+
+        """
+        if self.model_type == 'PHDNet':
+            return self.PHDNet_inference(composite_image, composite_mask)
+        if random_seed != None:
+            seed_everything(random_seed)
+        comp, mask = self.inputs_preprocess(composite_image, composite_mask)
+        self.scheduler.set_timesteps(sample_steps, device=self.device)
+        init_timestep = min(int(sample_steps * strength), sample_steps)
+        t_start       = max(sample_steps - init_timestep, 0)
+        timesteps     = self.scheduler.timesteps[t_start:]
+        
+        c = self.model.get_learned_conditioning([''])
+        batch_size = mask.shape[0]
+        latent_timestep = timesteps[:1].repeat(batch_size)
+
+        x_0 = self.model.encode_first_stage(comp)
+        x_0 = self.model.get_first_stage_encoding(x_0)
+        mask_latents = F.interpolate(mask, size=x_0.shape[-2:])
+        noise   = torch.randn(x_0.shape, device=x_0.device, dtype=x_0.dtype)
+        latents = self.scheduler.add_noise(x_0, noise, latent_timestep)
+
+        adapter_input = torch.cat((comp, mask), dim=1).to(dtype=comp.dtype)
+        features_adapter = self.adapter(adapter_input)
+
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            noise = torch.randn(x_0.shape, device=x_0.device, dtype=x_0.dtype)
+            t_latents = self.scheduler.add_noise(x_0, noise, t)
+            latents = latents * mask_latents + t_latents * (1 - mask_latents)
+
+            latent_model_input = latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            noise_pred = self.model.model.diffusion_model(x=latent_model_input,
+                                                          fg_mask=mask_latents, 
+                                                          timesteps=t.repeat(batch_size), 
+                                                          context=c, 
+                                                          features_adapter=features_adapter)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        x_samples_ddim = self.model.decode_first_stage(latents)
+        preds = self.outputs_postprocess(x_samples_ddim)
+        return preds
+