add all

configs/instant-mesh-base.yaml

@@ -0,0 +1,22 @@
+model_config:
+  target: src.models.lrm_mesh.InstantMesh
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 12
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 40
+    rendering_samples_per_ray: 96
+    grid_res: 128
+    grid_scale: 2.1
+
+
+infer_config:
+  unet_path: ckpts/diffusion_pytorch_model.bin
+  model_path: ckpts/instant_mesh_base.ckpt
+  texture_resolution: 1024
+  render_resolution: 512

configs/instant-mesh-large-train.yaml

@@ -0,0 +1,67 @@
+model:
+  base_learning_rate: 4.0e-05
+  target: src.model_mesh.MVRecon
+  params:
+    init_ckpt: logs/instant-nerf-large-train/checkpoints/last.ckpt
+    input_size: 320
+    render_size: 512
+    
+    lrm_generator_config:
+      target: src.models.lrm_mesh.InstantMesh
+      params:
+        encoder_feat_dim: 768
+        encoder_freeze: false
+        encoder_model_name: facebook/dino-vitb16
+        transformer_dim: 1024
+        transformer_layers: 16
+        transformer_heads: 16
+        triplane_low_res: 32
+        triplane_high_res: 64
+        triplane_dim: 80
+        rendering_samples_per_ray: 128
+        grid_res: 128
+        grid_scale: 2.1
+
+
+data:
+  target: src.data.objaverse.DataModuleFromConfig
+  params:
+    batch_size: 2
+    num_workers: 8
+    train:
+      target: src.data.objaverse.ObjaverseData
+      params:
+        root_dir: data/objaverse
+        meta_fname: filtered_obj_name.json
+        input_image_dir: rendering_random_32views
+        target_image_dir: rendering_random_32views
+        input_view_num: 6
+        target_view_num: 4
+        total_view_n: 32
+        fov: 50
+        camera_rotation: true
+        validation: false
+    validation:
+      target: src.data.objaverse.ValidationData
+      params:
+        root_dir: data/valid_samples
+        input_view_num: 6
+        input_image_size: 320
+        fov: 30
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 2000
+      save_top_k: -1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    val_check_interval: 1000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run

configs/instant-mesh-large.yaml

@@ -0,0 +1,22 @@
+model_config:
+  target: src.models.lrm_mesh.InstantMesh
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 16
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 80
+    rendering_samples_per_ray: 128
+    grid_res: 128
+    grid_scale: 2.1
+
+
+infer_config:
+  unet_path: ckpts/diffusion_pytorch_model.bin
+  model_path: ckpts/instant_mesh_large.ckpt
+  texture_resolution: 1024
+  render_resolution: 512

configs/instant-mesh-large_refine.yaml

@@ -0,0 +1,22 @@
+model_config:
+  target: src.models.lrm_mesh.InstantMesh
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 16
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 80
+    rendering_samples_per_ray: 128
+    grid_res: 128
+    grid_scale: 2.1
+
+
+infer_config:
+  unet_path: step=00260000.ckpt 
+  model_path: ckpts/instant_mesh_large.ckpt
+  texture_resolution: 8192
+  render_resolution: 1536

configs/instant-nerf-base.yaml

@@ -0,0 +1,21 @@
+model_config:
+  target: src.models.lrm.InstantNeRF
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 12
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 40
+    rendering_samples_per_ray: 96
+
+
+infer_config:
+  unet_path: ckpts/diffusion_pytorch_model.bin
+  model_path: ckpts/instant_nerf_base.ckpt
+  mesh_threshold: 10.0
+  mesh_resolution: 256
+  render_resolution: 384

configs/instant-nerf-large-best.yaml

@@ -0,0 +1,21 @@
+model_config:
+  target: src.models.lrm.InstantNeRF
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 16
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 80
+    rendering_samples_per_ray: 128
+
+
+infer_config:
+  unet_path: best_21.ckpt
+  model_path: ckpts/instant_nerf_large.ckpt
+  mesh_threshold: 5.0
+  mesh_resolution: 256
+  render_resolution: 512

configs/instant-nerf-large-train.yaml

@@ -0,0 +1,65 @@
+model:
+  base_learning_rate: 4.0e-04
+  target: src.model.MVRecon
+  params:
+    input_size: 320
+    render_size: 192
+    
+    lrm_generator_config:
+      target: src.models.lrm.InstantNeRF
+      params:
+        encoder_feat_dim: 768
+        encoder_freeze: false
+        encoder_model_name: facebook/dino-vitb16
+        transformer_dim: 1024
+        transformer_layers: 16
+        transformer_heads: 16
+        triplane_low_res: 32
+        triplane_high_res: 64
+        triplane_dim: 80
+        rendering_samples_per_ray: 128
+
+
+data:
+  target: src.data.objaverse.DataModuleFromConfig
+  params:
+    batch_size: 2
+    num_workers: 8
+    train:
+      target: src.data.objaverse.ObjaverseData
+      params:
+        root_dir: data/objaverse
+        meta_fname: filtered_obj_name.json
+        input_image_dir: rendering_random_32views
+        target_image_dir: rendering_random_32views
+        input_view_num: 6
+        target_view_num: 4
+        total_view_n: 32
+        fov: 50
+        camera_rotation: true
+        validation: false
+    validation:
+      target: src.data.objaverse.ValidationData
+      params:
+        root_dir: data/valid_samples
+        input_view_num: 6
+        input_image_size: 320
+        fov: 30
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1000
+      save_top_k: -1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 1000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run

configs/instant-nerf-large.yaml

@@ -0,0 +1,21 @@
+model_config:
+  target: src.models.lrm.InstantNeRF
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 16
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 80
+    rendering_samples_per_ray: 128
+
+
+infer_config:
+  unet_path: logs/zero123plus-refine_finetune_single_view/checkpoints/step=00210000.ckpt
+  model_path: ckpts/instant_nerf_large.ckpt
+  mesh_threshold: 10.0
+  mesh_resolution: 256
+  render_resolution: 320

configs/instant-nerf-sdedit.yaml

@@ -0,0 +1,21 @@
+model_config:
+  target: src.models.lrm.InstantNeRF
+  params:
+    encoder_feat_dim: 768
+    encoder_freeze: false
+    encoder_model_name: facebook/dino-vitb16
+    transformer_dim: 1024
+    transformer_layers: 16
+    transformer_heads: 16
+    triplane_low_res: 32
+    triplane_high_res: 64
+    triplane_dim: 80
+    rendering_samples_per_ray: 128
+
+
+infer_config:
+  unet_path: ckpts/diffusion_pytorch_model.bin
+  model_path: ckpts/instant_nerf_large.ckpt
+  mesh_threshold: 10.0
+  mesh_resolution: 256
+  render_resolution: 512

configs/zero123plus-finetune.yaml

@@ -0,0 +1,47 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusion
+  params:
+    drop_cond_prob: 0.1
+
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 1
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.ObjaverseData
+      params:
+        root_dir: data/objaverse
+        meta_fname: lvis-annotations.json
+        image_dir: rendering_zero123plus
+        validation: false
+    validation:
+      target: src.data.objaverse_zero123plus.ObjaverseData
+      params:
+        root_dir: data/objaverse
+        meta_fname: lvis-annotations.json
+        image_dir: rendering_zero123plus
+        validation: true
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1000
+      save_top_k: -1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 1000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run

configs/zero123plus-refine_finetune.yaml

@@ -0,0 +1,54 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusionRefinement
+  params:
+    drop_cond_prob: 0.1
+    refinement: true
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+    
+    
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 3
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: false
+        caption_path: captions.json
+        split_path: dataset_splits.json
+    validation:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: true
+        caption_path: captions.json
+        split_path: dataset_splits.json
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 10000
+      save_top_k: 1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 10000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run
+  

configs/zero123plus-refine_finetune_2.yaml

@@ -0,0 +1,51 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusionRefinement
+  params:
+    drop_cond_prob: 0.1
+    refinement: true
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+    
+    
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 3
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: false
+        overfit: true
+    validation:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: true
+        overfit: true
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1000
+      save_top_k: -1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 100
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run

configs/zero123plus-refine_finetune_relit.yaml

@@ -0,0 +1,52 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusionRefinement
+  params:
+    drop_cond_prob: 0.1
+    refinement: true
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+    
+    
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 3
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset_subset_relighted/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: false
+        caption_path: captions.json
+    validation:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset_subset_relighted/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: true
+        caption_path: captions.json
+
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1000
+      save_top_k: 1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 1000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run
+  

configs/zero123plus-refine_finetune_single_light.yaml

@@ -0,0 +1,56 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusionRefinement
+  params:
+    drop_cond_prob: 0.1
+    refinement: true
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+    
+    
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 3
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: false
+        caption_path: captions.json
+        split_path: dataset_splits_fixed.json
+        single_view: false
+        single_light: true
+    validation:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: true
+        caption_path: captions.json
+        split_path: dataset_splits_fixed.json
+        single_view: false
+        single_light: true
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 10000
+      save_top_k: 1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 10000000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run
+  

configs/zero123plus-refine_finetune_single_view.yaml

@@ -0,0 +1,55 @@
+model:
+  base_learning_rate: 1.0e-05
+  target: zero123plus.model.MVDiffusionRefinement
+  params:
+    drop_cond_prob: 0.1
+    refinement: true
+    stable_diffusion_config:
+      pretrained_model_name_or_path: sudo-ai/zero123plus-v1.2
+      custom_pipeline: ./zero123plus
+    
+    
+
+data:
+  target: src.data.objaverse_zero123plus.DataModuleFromConfig
+  params:
+    batch_size: 18
+    num_workers: 1
+    train:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: false
+        caption_path: captions.json
+        split_path: dataset_splits_fixed.json
+        single_view: true
+    validation:
+      target: src.data.objaverse_zero123plus.RefinementData
+      params:
+        root_dir: refinement_dataset/
+        gt_subpath: gt
+        pred_subpath: shap_e
+        validation: true
+        caption_path: captions.json
+        split_path: dataset_splits_fixed.json
+        single_view: true
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 10000
+      save_top_k: 1
+      save_last: true
+  callbacks: {}
+
+  trainer:
+    benchmark: true
+    max_epochs: -1
+    gradient_clip_val: 1.0
+    val_check_interval: 10000000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    check_val_every_n_epoch: null   # if not set this, validation does not run
+  

zero123plus/model.py

@@ -0,0 +1,547 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from tqdm import tqdm
+from torchvision.transforms import v2
+from torchvision.utils import make_grid, save_image
+from einops import rearrange
+
+from src.utils.train_util import instantiate_from_config
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, DDPMScheduler, UNet2DConditionModel
+from .pipeline import RefOnlyNoisedUNet
+
+
+def scale_latents(latents):
+    latents = (latents - 0.22) * 0.75
+    return latents
+
+
+def unscale_latents(latents):
+    latents = latents / 0.75 + 0.22
+    return latents
+
+
+def scale_image(image):
+    image = image * 0.5 / 0.8
+    return image
+
+
+def unscale_image(image):
+    image = image / 0.5 * 0.8
+    return image
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+class MVDiffusionRefinement(pl.LightningModule):
+    def __init__(
+        self,
+        stable_diffusion_config,
+        refinement,
+        drop_cond_prob=0.1,
+    ):
+        super(MVDiffusionRefinement, self).__init__()
+
+        self.drop_cond_prob = drop_cond_prob
+        self.refinement = refinement
+        self.register_schedule()
+
+        # init modules
+
+        pipeline = DiffusionPipeline.from_pretrained(**stable_diffusion_config,low_cpu_mem_usage=False)
+        
+        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+            pipeline.scheduler.config, timestep_spacing='trailing'
+        )
+
+        self.pipeline = pipeline
+        if refinement:
+            from huggingface_hub import hf_hub_download
+            unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model")
+            state_dict = torch.load(unet_ckpt_path, map_location='cpu')
+            self.pipeline.unet.load_state_dict(state_dict, strict=False)
+            pipeline.unet.load_state_dict(state_dict, strict=False)
+        train_sched = DDPMScheduler.from_config(self.pipeline.scheduler.config)
+
+        in_channels = 8
+        out_channels = self.pipeline.unet.conv_in.out_channels
+        self.pipeline.unet.register_to_config(in_channels=in_channels)
+        with torch.no_grad():
+            new_conv_in = nn.Conv2d(
+                in_channels, out_channels, self.pipeline.unet.conv_in.kernel_size, self.pipeline.unet.conv_in.stride, self.pipeline.unet.conv_in.padding
+            )
+            new_conv_in.weight.zero_()
+            new_conv_in.weight[:, :4, :, :].copy_(self.pipeline.unet.conv_in.weight)
+            self.pipeline.unet.conv_in = new_conv_in
+
+        if isinstance(self.pipeline.unet, UNet2DConditionModel):
+            self.pipeline.unet = RefOnlyNoisedUNet(self.pipeline.unet, train_sched, self.pipeline.scheduler)
+        
+
+        self.train_scheduler = train_sched      # use ddpm scheduler during training
+
+        self.unet = pipeline.unet
+
+        # validation output buffer
+        self.validation_step_outputs = []
+        with torch.no_grad():
+            self.cond_latents_zero = self.encode_condition_image(torch.zeros(1,3,320,320)).to(self.device)
+            self.prompt_latents_zero = self.pipeline._encode_prompt([""], self.device, 1, False)
+
+
+    def register_schedule(self):
+        self.num_timesteps = 1000
+
+        # replace scaled_linear schedule with linear schedule as Zero123++
+        beta_start = 0.00085
+        beta_end = 0.0120
+        betas = torch.linspace(beta_start, beta_end, 1000, dtype=torch.float32)
+        
+        alphas = 1. - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0)
+
+        self.register_buffer('betas', betas.float())
+        self.register_buffer('alphas_cumprod', alphas_cumprod.float())
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float())
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod).float())
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod).float())
+        
+        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod).float())
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1).float())
+    
+    def on_fit_start(self):
+        device = torch.device(f'cuda:{self.global_rank}')
+        self.pipeline.to(device)
+        if self.global_rank == 0:
+            os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)
+            os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)
+    
+    def prepare_batch_data(self, batch):
+        unrefined_imgs = batch['unrefined_imgs']  # (B, 6, C, H, W)
+        unrefined_imgs = v2.functional.resize(unrefined_imgs, 320, interpolation=3, antialias=True).clamp(0, 1)
+        unrefined_imgs = rearrange(unrefined_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2)    # (B, C, 3H, 2W)
+        unrefined_imgs = unrefined_imgs.to(self.device)
+        
+        target_imgs = batch['refined_imgs']  # (B, 6, C, H, W)
+        target_imgs = v2.functional.resize(target_imgs, 320, interpolation=3, antialias=True).clamp(0, 1)
+        target_imgs = rearrange(target_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2)    # (B, C, 3H, 2W)
+        target_imgs = target_imgs.to(self.device)
+        return unrefined_imgs, target_imgs
+           
+    
+    @torch.no_grad()
+    def forward_vision_encoder(self, images):
+        dtype = next(self.pipeline.vision_encoder.parameters()).dtype
+        image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
+        image_pt = self.pipeline.feature_extractor_clip(images=image_pil, return_tensors="pt").pixel_values
+        image_pt = image_pt.to(device=self.device, dtype=dtype)
+        global_embeds = self.pipeline.vision_encoder(image_pt, output_hidden_states=False).image_embeds
+        global_embeds = global_embeds.unsqueeze(-2)
+
+        encoder_hidden_states = self.pipeline._encode_prompt("", self.device, 1, False)[0]
+        ramp = global_embeds.new_tensor(self.pipeline.config.ramping_coefficients).unsqueeze(-1)
+        encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+
+        return encoder_hidden_states
+    
+    @torch.no_grad()
+    def encode_condition_image(self, images):
+        dtype = next(self.pipeline.vae.parameters()).dtype
+        image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
+        image_pt = self.pipeline.feature_extractor_vae(images=image_pil, return_tensors="pt").pixel_values
+        image_pt = image_pt.to(device=self.device, dtype=dtype)
+        latents = self.pipeline.vae.encode(image_pt).latent_dist.sample()
+        return latents
+    
+    @torch.no_grad()
+    def encode_target_images(self, images):
+        dtype = next(self.pipeline.vae.parameters()).dtype
+        # equals to scaling images to [-1, 1] first and then call scale_image
+        
+        images = (images - 0.5) / 0.8   # [-0.625, 0.625]
+        posterior = self.pipeline.vae.encode(images.to(dtype)).latent_dist
+        latents = posterior.sample() * self.pipeline.vae.config.scaling_factor
+        latents = scale_latents(latents)
+        return latents
+    
+    def forward_unet(self, latents, t, prompt_embeds, cond_latents, cross_attention_kwargs=None):
+        dtype = next(self.pipeline.unet.parameters()).dtype
+        latents = latents.to(dtype)
+        prompt_embeds = prompt_embeds.to(dtype)
+        cond_latents = cond_latents.to(dtype)
+        if cross_attention_kwargs is None:
+            cross_attention_kwargs = dict()
+        cross_attention_kwargs.update(cond_lat=cond_latents)
+        # cross_attention_kwargs = dict(cond_lat=cond_latents)
+        
+        pred_noise = self.pipeline.unet(
+            latents,
+            t,
+            encoder_hidden_states=prompt_embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            return_dict=False,
+        )[0]
+        return pred_noise
+    
+    def predict_start_from_z_and_v(self, x_t, t, v):
+        return (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
+            extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
+        )
+
+    def get_v(self, x, noise, t):
+        return (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
+            extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
+        )
+        
+    def decode_latents(self, latents_pred):
+        latents = unscale_latents(latents_pred)
+        images = unscale_image(self.pipeline.vae.decode(latents / self.pipeline.vae.config.scaling_factor, return_dict=False)[0])   # [-1, 1]
+        images = (images * 0.5 + 0.5).clamp(0, 1)
+        return images
+
+    def training_step(self, batch, batch_idx):
+        # get input
+        latents_source, latents_target = batch['unrefined_imgs'], batch['refined_imgs']
+        captions = batch['caption']
+        # sample random timestep
+        B = latents_source.shape[0]
+        
+        t = torch.randint(0, self.num_timesteps, size=(B,)).long().to(self.device)
+
+        # classifier-free guidance
+        if np.random.rand() < self.drop_cond_prob:
+            prompt_embeds = self.prompt_latents_zero.to(self.device).expand(B, -1, -1)
+        else:
+            prompt_embeds = self.pipeline._encode_prompt(captions,self.device, 1, False)
+        cond_latents = self.cond_latents_zero.to(self.device)
+
+        # with torch.no_grad():
+            # latents_source = self.pipeline.vae.encode(source_imgs).latent_dist.mode()
+        noise = torch.randn_like(latents_target)
+        latents_noisy = self.train_scheduler.add_noise(latents_target, noise, t)
+        latents_noisy_unet = torch.cat([latents_noisy, latents_source], dim=1)
+        cak = dict(dont_forward_cond_state=True)
+        v_pred = self.forward_unet(latents_noisy_unet, t, prompt_embeds, cond_latents, cross_attention_kwargs=cak)
+        v_target = self.get_v(latents_target, noise, t)
+
+        loss, loss_dict = self.compute_loss(v_pred, v_target)
+
+        # logging
+        self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+        self.log("global_step", self.global_step, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+        lr = self.optimizers().param_groups[0]['lr']
+        self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        if self.global_step % 5000000 == 0 and self.global_rank == 0:
+            with torch.no_grad():
+                latents_pred = self.predict_start_from_z_and_v(latents_noisy, t, v_pred)
+                images = self.decode_latents(latents_pred)
+                target_imgs = self.decode_latents(latents_target)
+
+                images = torch.cat([target_imgs, images], dim=-2)
+
+                grid = make_grid(images, nrow=images.shape[0], normalize=True, value_range=(0, 1))
+                save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png'))
+
+        return loss
+        
+    def compute_loss(self, noise_pred, noise_gt):
+        loss = F.mse_loss(noise_pred, noise_gt)
+
+        prefix = 'train'
+        loss_dict = {}
+        loss_dict.update({f'{prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        # get input
+        latents_source, latents_target = batch['unrefined_imgs'], batch['refined_imgs']
+        prompts = batch['caption']
+        source_imgs = self.decode_latents(latents_source)
+        target_imgs = self.decode_latents(latents_target)
+
+        images_pil = [v2.functional.to_pil_image(source_imgs[i]) for i in range(source_imgs.shape[0])]
+
+        outputs = []
+        for source_img,prompt in zip(images_pil,prompts):
+            latent = self.pipeline.refine(source_img,prompt=prompt, num_inference_steps=75, output_type='latent').images
+            image = unscale_image(self.pipeline.vae.decode(latent / self.pipeline.vae.config.scaling_factor, return_dict=False)[0])  # [-1, 1]
+            image = (image * 0.5 + 0.5).clamp(0, 1)
+            outputs.append(image)
+        outputs = torch.cat(outputs, dim=0).to(self.device)
+        images = torch.cat([target_imgs, outputs, source_imgs], dim=-2)
+        
+        self.validation_step_outputs.append(images)
+    
+    @torch.no_grad()
+    def on_validation_epoch_end(self):
+        images = torch.cat(self.validation_step_outputs, dim=0)
+        all_images = self.all_gather(images)
+        # all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')
+        imgs = all_images.chunk(all_images.shape[0], dim=0)
+
+        if self.global_rank == 0:
+            os.makedirs(os.path.join(self.logdir, 'images_val', f'{self.global_step:07d}'), exist_ok=True)
+            grid = make_grid(all_images, nrow=8, normalize=True, value_range=(0, 1))
+            save_image(grid, os.path.join(self.logdir, 'images_val',f'{self.global_step:07d}', f'all.png'))
+            for idx, img in enumerate(imgs):
+                target, output, source = img.chunk(3, dim=-2)
+                img = torch.cat([source, target, output], dim=-1)
+                save_image(img, os.path.join(self.logdir, 'images_val',f'{self.global_step:07d}', f'comparison_img_{idx}.png'))
+                source_outputs = torch.cat([source, output], dim=-1)
+                save_image(source_outputs, os.path.join(self.logdir, 'images_val',f'{self.global_step:07d}', f'comparison_source_output_img_{idx}.png'))
+        self.validation_step_outputs.clear()  # free memory
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+
+        optimizer = torch.optim.AdamW(self.unet.parameters(), lr=lr)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3000, eta_min=lr/4)
+
+        return {'optimizer': optimizer, 'lr_scheduler': scheduler}
+
+class MVDiffusion(pl.LightningModule):
+    def __init__(
+        self,
+        stable_diffusion_config,
+        drop_cond_prob=0.2,
+    ):
+        super(MVDiffusion, self).__init__()
+
+        self.drop_cond_prob = drop_cond_prob
+        self.register_schedule()
+
+        # init modules
+
+        pipeline = DiffusionPipeline.from_pretrained(**stable_diffusion_config)
+        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+            pipeline.scheduler.config, timestep_spacing='trailing'
+        )
+
+        self.pipeline = pipeline
+        train_sched = DDPMScheduler.from_config(self.pipeline.scheduler.config)
+        if isinstance(self.pipeline.unet, UNet2DConditionModel):
+            self.pipeline.unet = RefOnlyNoisedUNet(self.pipeline.unet, train_sched, self.pipeline.scheduler)
+        
+
+        self.train_scheduler = train_sched      # use ddpm scheduler during training
+
+        self.unet = pipeline.unet
+
+        # validation output buffer
+        self.validation_step_outputs = []
+
+    def register_schedule(self):
+        self.num_timesteps = 1000
+
+        # replace scaled_linear schedule with linear schedule as Zero123++
+        beta_start = 0.00085
+        beta_end = 0.0120
+        betas = torch.linspace(beta_start, beta_end, 1000, dtype=torch.float32)
+        
+        alphas = 1. - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0)
+
+        self.register_buffer('betas', betas.float())
+        self.register_buffer('alphas_cumprod', alphas_cumprod.float())
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float())
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod).float())
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod).float())
+        
+        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod).float())
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1).float())
+    
+    def on_fit_start(self):
+        device = torch.device(f'cuda:{self.global_rank}')
+        self.pipeline.to(device)
+        if self.global_rank == 0:
+            os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)
+            os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)
+    
+    
+    def prepare_batch_data(self, batch):
+        cond_imgs = batch['cond_imgs']      # (B, C, H, W)
+        cond_imgs = cond_imgs.to(self.device)
+
+        # random resize the condition image
+        cond_size = np.random.randint(128, 513)
+        cond_imgs = v2.functional.resize(cond_imgs, cond_size, interpolation=3, antialias=True).clamp(0, 1)
+
+        target_imgs = batch['target_imgs']  # (B, 6, C, H, W)
+        target_imgs = v2.functional.resize(target_imgs, 320, interpolation=3, antialias=True).clamp(0, 1)
+        target_imgs = rearrange(target_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2)    # (B, C, 3H, 2W)
+        target_imgs = target_imgs.to(self.device)
+
+        return cond_imgs, target_imgs
+           
+    
+    @torch.no_grad()
+    def forward_vision_encoder(self, images):
+        dtype = next(self.pipeline.vision_encoder.parameters()).dtype
+        image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
+        image_pt = self.pipeline.feature_extractor_clip(images=image_pil, return_tensors="pt").pixel_values
+        image_pt = image_pt.to(device=self.device, dtype=dtype)
+        global_embeds = self.pipeline.vision_encoder(image_pt, output_hidden_states=False).image_embeds
+        global_embeds = global_embeds.unsqueeze(-2)
+
+        encoder_hidden_states = self.pipeline._encode_prompt("", self.device, 1, False)[0]
+        ramp = global_embeds.new_tensor(self.pipeline.config.ramping_coefficients).unsqueeze(-1)
+        encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+
+        return encoder_hidden_states
+    
+    @torch.no_grad()
+    def encode_condition_image(self, images):
+        dtype = next(self.pipeline.vae.parameters()).dtype
+        image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
+        image_pt = self.pipeline.feature_extractor_vae(images=image_pil, return_tensors="pt").pixel_values
+        image_pt = image_pt.to(device=self.device, dtype=dtype)
+        latents = self.pipeline.vae.encode(image_pt).latent_dist.sample()
+        return latents
+    
+    @torch.no_grad()
+    def encode_target_images(self, images):
+        dtype = next(self.pipeline.vae.parameters()).dtype
+        # equals to scaling images to [-1, 1] first and then call scale_image
+        images = (images - 0.5) / 0.8   # [-0.625, 0.625]
+        posterior = self.pipeline.vae.encode(images.to(dtype)).latent_dist
+        latents = posterior.sample() * self.pipeline.vae.config.scaling_factor
+        latents = scale_latents(latents)
+        return latents
+    
+    def forward_unet(self, latents, t, prompt_embeds, cond_latents):
+        dtype = next(self.pipeline.unet.parameters()).dtype
+        latents = latents.to(dtype)
+        prompt_embeds = prompt_embeds.to(dtype)
+        cond_latents = cond_latents.to(dtype)
+        cross_attention_kwargs = dict(cond_lat=cond_latents)
+        pred_noise = self.pipeline.unet(
+            latents,
+            t,
+            encoder_hidden_states=prompt_embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            return_dict=False,
+        )[0]
+        return pred_noise
+    
+    def predict_start_from_z_and_v(self, x_t, t, v):
+        return (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
+            extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
+        )
+
+    def get_v(self, x, noise, t):
+        return (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
+            extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
+        )
+    
+    def training_step(self, batch, batch_idx):
+        # get input
+        cond_imgs, target_imgs = self.prepare_batch_data(batch)
+
+        # sample random timestep
+        B = cond_imgs.shape[0]
+        
+        t = torch.randint(0, self.num_timesteps, size=(B,)).long().to(self.device)
+
+        # classifier-free guidance
+        if np.random.rand() < self.drop_cond_prob:
+            prompt_embeds = self.pipeline._encode_prompt([""]*B, self.device, 1, False)
+            cond_latents = self.encode_condition_image(torch.zeros_like(cond_imgs))
+        else:
+            prompt_embeds = self.forward_vision_encoder(cond_imgs)
+            cond_latents = self.encode_condition_image(cond_imgs)
+
+        latents = self.encode_target_images(target_imgs)
+        noise = torch.randn_like(latents)
+        latents_noisy = self.train_scheduler.add_noise(latents, noise, t)
+        
+        v_pred = self.forward_unet(latents_noisy, t, prompt_embeds, cond_latents)
+        v_target = self.get_v(latents, noise, t)
+
+        loss, loss_dict = self.compute_loss(v_pred, v_target)
+
+        # logging
+        self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+        self.log("global_step", self.global_step, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+        lr = self.optimizers().param_groups[0]['lr']
+        self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        if self.global_step % 50 == 0 and self.global_rank == 0:
+            with torch.no_grad():
+                latents_pred = self.predict_start_from_z_and_v(latents_noisy, t, v_pred)
+
+                latents = unscale_latents(latents_pred)
+                images = unscale_image(self.pipeline.vae.decode(latents / self.pipeline.vae.config.scaling_factor, return_dict=False)[0])   # [-1, 1]
+                images = (images * 0.5 + 0.5).clamp(0, 1)
+                images = torch.cat([target_imgs, images], dim=-2)
+
+                grid = make_grid(images, nrow=images.shape[0], normalize=True, value_range=(0, 1))
+                save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png'))
+
+        return loss
+        
+    def compute_loss(self, noise_pred, noise_gt):
+        loss = F.mse_loss(noise_pred, noise_gt)
+
+        prefix = 'train'
+        loss_dict = {}
+        loss_dict.update({f'{prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        # get input
+        cond_imgs, target_imgs = self.prepare_batch_data(batch)
+
+        images_pil = [v2.functional.to_pil_image(cond_imgs[i]) for i in range(cond_imgs.shape[0])]
+
+        outputs = []
+        for cond_img in images_pil:
+            latent = self.pipeline(cond_img, num_inference_steps=75, output_type='latent').images
+            image = unscale_image(self.pipeline.vae.decode(latent / self.pipeline.vae.config.scaling_factor, return_dict=False)[0])   # [-1, 1]
+            image = (image * 0.5 + 0.5).clamp(0, 1)
+            outputs.append(image)
+        outputs = torch.cat(outputs, dim=0).to(self.device)
+        images = torch.cat([target_imgs, outputs], dim=-2)
+        
+        self.validation_step_outputs.append(images)
+    
+    @torch.no_grad()
+    def on_validation_epoch_end(self):
+        images = torch.cat(self.validation_step_outputs, dim=0)
+
+        all_images = self.all_gather(images)
+        all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')
+
+        if self.global_rank == 0:
+            grid = make_grid(all_images, nrow=8, normalize=True, value_range=(0, 1))
+            save_image(grid, os.path.join(self.logdir, 'images_val', f'val_{self.global_step:07d}.png'))
+
+        self.validation_step_outputs.clear()  # free memory
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+
+        optimizer = torch.optim.AdamW(self.unet.parameters(), lr=lr)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3000, eta_min=lr/4)
+
+        return {'optimizer': optimizer, 'lr_scheduler': scheduler}

zero123plus/pipeline.py

@@ -0,0 +1,1125 @@
+from typing import Any, Dict, Optional
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.schedulers import KarrasDiffusionSchedulers
+
+import numpy
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+import torch.distributed
+import transformers
+from collections import OrderedDict
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.utils import randn_tensor
+import diffusers
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DiffusionPipeline,
+    EulerAncestralDiscreteScheduler,
+    UNet2DConditionModel,
+    ImagePipelineOutput,
+)
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.attention_processor import (
+    Attention,
+    AttnProcessor,
+    XFormersAttnProcessor,
+    AttnProcessor2_0,
+)
+from diffusers.utils.import_utils import is_xformers_available
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def to_rgb_image(maybe_rgba: Image.Image):
+    if maybe_rgba.mode == "RGB":
+        return maybe_rgba
+    elif maybe_rgba.mode == "RGBA":
+        rgba = maybe_rgba
+        img = numpy.random.randint(
+            255, 256, size=[rgba.size[1], rgba.size[0], 3], dtype=numpy.uint8
+        )
+        img = Image.fromarray(img, "RGB")
+        img.paste(rgba, mask=rgba.getchannel("A"))
+        return img
+    else:
+        raise ValueError("Unsupported image type.", maybe_rgba.mode)
+
+
+class ReferenceOnlyAttnProc(torch.nn.Module):
+    def __init__(self, chained_proc, enabled=False, name=None) -> None:
+        super().__init__()
+        self.enabled = enabled
+        self.chained_proc = chained_proc
+        self.name = name
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        mode="w",
+        ref_dict: dict = None,
+        is_cfg_guidance=False,
+    ) -> Any:
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        if self.enabled and is_cfg_guidance:
+            res0 = self.chained_proc(
+                attn, hidden_states[:1], encoder_hidden_states[:1], attention_mask
+            )
+            hidden_states = hidden_states[1:]
+            encoder_hidden_states = encoder_hidden_states[1:]
+        if self.enabled:
+            if mode == "w":
+                ref_dict[self.name] = encoder_hidden_states
+            elif mode == "r":
+                encoder_hidden_states = torch.cat(
+                    [encoder_hidden_states, ref_dict.pop(self.name)], dim=1
+                )
+            elif mode == "m":
+                encoder_hidden_states = torch.cat(
+                    [encoder_hidden_states, ref_dict[self.name]], dim=1
+                )
+            elif mode == "c":
+                encoder_hidden_states = torch.cat(
+                    [encoder_hidden_states, encoder_hidden_states], dim=1
+                )
+            else:
+                assert False, mode
+        res = self.chained_proc(
+            attn, hidden_states, encoder_hidden_states, attention_mask
+        )
+        if self.enabled and is_cfg_guidance:
+            res = torch.cat([res0, res])
+        return res
+
+
+class RefOnlyNoisedUNet(torch.nn.Module):
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        train_sched: DDPMScheduler,
+        val_sched: EulerAncestralDiscreteScheduler,
+    ) -> None:
+        super().__init__()
+        self.unet = unet
+        self.train_sched = train_sched
+        self.val_sched = val_sched
+
+        unet_lora_attn_procs = dict()
+        for name, _ in unet.attn_processors.items():
+            if torch.__version__ >= "2.0":
+                default_attn_proc = AttnProcessor2_0()
+            elif is_xformers_available():
+                default_attn_proc = XFormersAttnProcessor()
+            else:
+                default_attn_proc = AttnProcessor()
+            unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(
+                default_attn_proc, enabled=name.endswith("attn1.processor"), name=name
+            )
+        unet.set_attn_processor(unet_lora_attn_procs)
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.unet, name)
+
+    def forward_cond(
+        self,
+        noisy_cond_lat,
+        timestep,
+        encoder_hidden_states,
+        class_labels,
+        ref_dict,
+        is_cfg_guidance,
+        **kwargs,
+    ):
+        if is_cfg_guidance:
+            encoder_hidden_states = encoder_hidden_states[1:]
+            class_labels = class_labels[1:]
+        self.unet(
+            noisy_cond_lat,
+            timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            class_labels=class_labels,
+            cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict),
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        sample,
+        timestep,
+        encoder_hidden_states,
+        class_labels=None,
+        *args,
+        cross_attention_kwargs,
+        down_block_res_samples=None,
+        mid_block_res_sample=None,
+        forward_cond_state=True,
+        **kwargs,
+    ):
+        cond_lat = cross_attention_kwargs["cond_lat"]
+        is_cfg_guidance = cross_attention_kwargs.get("is_cfg_guidance", False)
+        noise = torch.randn_like(cond_lat)
+        if self.training:
+            noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)
+            noisy_cond_lat = self.train_sched.scale_model_input(
+                noisy_cond_lat, timestep
+            )
+        else:
+            noisy_cond_lat = self.val_sched.add_noise(
+                cond_lat, noise, timestep.reshape(-1)
+            )
+            noisy_cond_lat = self.val_sched.scale_model_input(
+                noisy_cond_lat, timestep.reshape(-1)
+            )
+        ref_dict = {}
+        if "dont_forward_cond_state" not in cross_attention_kwargs.keys():
+            self.forward_cond(
+                noisy_cond_lat,
+                timestep,
+                encoder_hidden_states,
+                class_labels,
+                ref_dict,
+                is_cfg_guidance,
+                **kwargs,
+            )
+            mode = "r"
+        else:
+            mode = "c"
+        weight_dtype = self.unet.dtype
+        return self.unet(
+            sample,
+            timestep,
+            encoder_hidden_states,
+            *args,
+            class_labels=class_labels,
+            cross_attention_kwargs=dict(
+                mode=mode, ref_dict=ref_dict, is_cfg_guidance=is_cfg_guidance
+            ),
+            down_block_additional_residuals=[
+                sample.to(dtype=weight_dtype) for sample in down_block_res_samples
+            ]
+            if down_block_res_samples is not None
+            else None,
+            mid_block_additional_residual=(
+                mid_block_res_sample.to(dtype=weight_dtype)
+                if mid_block_res_sample is not None
+                else None
+            ),
+            **kwargs,
+        )
+
+
+def scale_latents(latents):
+    latents = (latents - 0.22) * 0.75
+    return latents
+
+
+def unscale_latents(latents):
+    latents = latents / 0.75 + 0.22
+    return latents
+
+
+def scale_image(image):
+    image = image * 0.5 / 0.8
+    return image
+
+
+def unscale_image(image):
+    image = image / 0.5 * 0.8
+    return image
+
+
+class DepthControlUNet(torch.nn.Module):
+    def __init__(
+        self,
+        unet: RefOnlyNoisedUNet,
+        controlnet: Optional[diffusers.ControlNetModel] = None,
+        conditioning_scale=1.0,
+    ) -> None:
+        super().__init__()
+        self.unet = unet
+        if controlnet is None:
+            self.controlnet = diffusers.ControlNetModel.from_unet(unet.unet)
+        else:
+            self.controlnet = controlnet
+        DefaultAttnProc = AttnProcessor2_0
+        if is_xformers_available():
+            DefaultAttnProc = XFormersAttnProcessor
+        self.controlnet.set_attn_processor(DefaultAttnProc())
+        self.conditioning_scale = conditioning_scale
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.unet, name)
+
+    def forward(
+        self,
+        sample,
+        timestep,
+        encoder_hidden_states,
+        class_labels=None,
+        *args,
+        cross_attention_kwargs: dict,
+        **kwargs,
+    ):
+        cross_attention_kwargs = dict(cross_attention_kwargs)
+        control_depth = cross_attention_kwargs.pop("control_depth")
+        down_block_res_samples, mid_block_res_sample = self.controlnet(
+            sample,
+            timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            controlnet_cond=control_depth,
+            conditioning_scale=self.conditioning_scale,
+            return_dict=False,
+        )
+        return self.unet(
+            sample,
+            timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            down_block_res_samples=down_block_res_samples,
+            mid_block_res_sample=mid_block_res_sample,
+            cross_attention_kwargs=cross_attention_kwargs,
+        )
+
+
+class ModuleListDict(torch.nn.Module):
+    def __init__(self, procs: dict) -> None:
+        super().__init__()
+        self.keys = sorted(procs.keys())
+        self.values = torch.nn.ModuleList(procs[k] for k in self.keys)
+
+    def __getitem__(self, key):
+        return self.values[self.keys.index(key)]
+
+
+class SuperNet(torch.nn.Module):
+    def __init__(self, state_dict: Dict[str, torch.Tensor]):
+        super().__init__()
+        state_dict = OrderedDict((k, state_dict[k]) for k in sorted(state_dict.keys()))
+        self.layers = torch.nn.ModuleList(state_dict.values())
+        self.mapping = dict(enumerate(state_dict.keys()))
+        self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}
+
+        # .processor for unet, .self_attn for text encoder
+        self.split_keys = [".processor", ".self_attn"]
+
+        # we add a hook to state_dict() and load_state_dict() so that the
+        # naming fits with `unet.attn_processors`
+        def map_to(module, state_dict, *args, **kwargs):
+            new_state_dict = {}
+            for key, value in state_dict.items():
+                num = int(key.split(".")[1])  # 0 is always "layers"
+                new_key = key.replace(f"layers.{num}", module.mapping[num])
+                new_state_dict[new_key] = value
+
+            return new_state_dict
+
+        def remap_key(key, state_dict):
+            for k in self.split_keys:
+                if k in key:
+                    return key.split(k)[0] + k
+            return key.split(".")[0]
+
+        def map_from(module, state_dict, *args, **kwargs):
+            all_keys = list(state_dict.keys())
+            for key in all_keys:
+                replace_key = remap_key(key, state_dict)
+                new_key = key.replace(
+                    replace_key, f"layers.{module.rev_mapping[replace_key]}"
+                )
+                state_dict[new_key] = state_dict[key]
+                del state_dict[key]
+
+        self._register_state_dict_hook(map_to)
+        self._register_load_state_dict_pre_hook(map_from, with_module=True)
+
+
+class Zero123PlusPipeline(diffusers.StableDiffusionPipeline):
+    tokenizer: transformers.CLIPTokenizer
+    text_encoder: transformers.CLIPTextModel
+    vision_encoder: transformers.CLIPVisionModelWithProjection
+
+    feature_extractor_clip: transformers.CLIPImageProcessor
+    unet: UNet2DConditionModel
+    scheduler: diffusers.schedulers.KarrasDiffusionSchedulers
+
+    vae: AutoencoderKL
+    ramping: nn.Linear
+
+    feature_extractor_vae: transformers.CLIPImageProcessor
+
+    depth_transforms_multi = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
+    )
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        vision_encoder: transformers.CLIPVisionModelWithProjection,
+        feature_extractor_clip: CLIPImageProcessor,
+        feature_extractor_vae: CLIPImageProcessor,
+        ramping_coefficients: Optional[list] = None,
+        safety_checker=None,
+    ):
+        DiffusionPipeline.__init__(self)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=None,
+            vision_encoder=vision_encoder,
+            feature_extractor_clip=feature_extractor_clip,
+            feature_extractor_vae=feature_extractor_vae,
+        )
+        self.register_to_config(ramping_coefficients=ramping_coefficients)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def prepare(self):
+        train_sched = DDPMScheduler.from_config(self.scheduler.config)
+        if isinstance(self.unet, UNet2DConditionModel):
+            self.unet = RefOnlyNoisedUNet(self.unet, train_sched, self.scheduler).eval()
+
+    def add_controlnet(
+        self,
+        controlnet: Optional[diffusers.ControlNetModel] = None,
+        conditioning_scale=1.0,
+    ):
+        self.prepare()
+        self.unet = DepthControlUNet(self.unet, controlnet, conditioning_scale)
+        return SuperNet(OrderedDict([("controlnet", self.unet.controlnet)]))
+
+    def encode_condition_image(self, image: torch.Tensor):
+        image = self.vae.encode(image).latent_dist.sample()
+        return image
+
+    @torch.no_grad()
+    def edit_latents(
+        self,
+        image_guidance: Image.Image,
+        multiview_source_image: Image.Image = None,
+        edit_strength: float = 1.0,
+        prompt="",
+        *args,
+        guidance_scale=0.0,
+        output_type: Optional[str] = "pil",
+        width=640,
+        height=960,
+        num_inference_steps=28,
+        return_dict=True,
+        **kwargs,
+    ):
+        self.prepare()
+        if image_guidance is None:
+            raise ValueError(
+                "Inputting embeddings not supported for this pipeline. Please pass an image."
+            )
+        if multiview_source_image is None:
+            raise ValueError("Multiview source image is required for this pipeline.")
+        assert not isinstance(image_guidance, torch.Tensor)
+        assert not isinstance(multiview_source_image, torch.Tensor)
+        image_guidance = to_rgb_image(image_guidance)
+        image_source = to_rgb_image(multiview_source_image)
+        image_guidance_1 = self.feature_extractor_vae(
+            images=image_guidance, return_tensors="pt"
+        ).pixel_values
+        image_guidance_2 = self.feature_extractor_clip(
+            images=image_source, return_tensors="pt"
+        ).pixel_values
+        image_guidance = image_guidance_1.to(
+            device=self.vae.device, dtype=self.vae.dtype
+        )
+        image_guidance_2 = image_guidance_2.to(
+            device=self.vae.device, dtype=self.vae.dtype
+        )
+
+        cond_lat = self.encode_condition_image(image_guidance)
+        # if guidance_scale > 1:
+        negative_lat = self.encode_condition_image(torch.zeros_like(image_guidance))
+        cond_lat = torch.cat([negative_lat, cond_lat])
+        encoded = self.vision_encoder(image_guidance_2, output_hidden_states=False)
+
+        global_embeds = encoded.image_embeds
+        global_embeds = global_embeds.unsqueeze(-2)
+        if hasattr(self, "encode_prompt"):
+            encoder_hidden_states = self.encode_prompt(prompt, self.device, 1, False)[0]
+        else:
+            encoder_hidden_states = self._encode_prompt(prompt, self.device, 1, False)
+        ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+        cak = dict(cond_lat=cond_lat)
+        mv_image = (
+            torch.from_numpy(numpy.array(multiview_source_image)).to(self.vae.device)
+            / 255.0
+        )
+        mv_image = (
+            mv_image.permute(2, 0, 1)
+            .to(self.vae.device)
+            .to(self.vae.dtype)
+            .unsqueeze(0)
+        )
+        latents = (
+            self.vae.encode(mv_image * 2.0 - 1.0).latent_dist.sample()
+            * self.vae.config.scaling_factor
+        )
+        latents: torch.Tensor = (
+            super()
+            .__call__(
+                None,
+                *args,
+                cross_attention_kwargs=cak,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=1,
+                prompt_embeds=encoder_hidden_states,
+                num_inference_steps=num_inference_steps,
+                output_type="latent",
+                width=width,
+                height=height,
+                latents=latents,
+                edit_strength=edit_strength,
+                **kwargs,
+            )
+            .images
+        )
+        latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = unscale_image(
+                self.vae.decode(
+                    latents / self.vae.config.scaling_factor, return_dict=False
+                )[0]
+            )
+        else:
+            image = latents
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+    @torch.no_grad()
+    def encode_target_images(self, images):
+        dtype = next(self.vae.parameters()).dtype
+        # equals to scaling images to [-1, 1] first and then call scale_image
+        images = (images - 0.5) / 0.8  # [-0.625, 0.625]
+        posterior = self.vae.encode(images.to(dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        latents = scale_latents(latents)
+        return latents
+
+    @torch.no_grad()
+    def sdedit(
+        self,
+        image,
+        *args,
+        cond_image: Image.Image = None,
+        output_type: Optional[str] = "pil",
+        width=640,
+        height=960,
+        num_inference_steps=75,
+        edit_strength=1.0,
+        return_dict=True,
+        guidance_scale=0.0,
+        **kwargs,
+    ):
+        self.prepare()
+        if image is None:
+            raise ValueError(
+                "Inputting embeddings not supported for this pipeline. Please pass an image."
+            )
+        assert not isinstance(image, torch.Tensor)
+        image = to_rgb_image(image)
+
+        # cond_lat = self.encode_condition_image(image_guidance)
+        if hasattr(self, "encode_prompt"):
+            encoder_hidden_states = self.encode_prompt([""], self.device, 1, False)[0]
+        else:
+            encoder_hidden_states = self._encode_prompt([""], self.device, 1, False)
+        # negative_lat = self.encode_condition_image(torch.zeros_like(image_guidance))
+        # cond_lat = torch.cat([negative_lat, cond_lat])
+        # encoded = self.vision_encoder(image_guidance_2, output_hidden_states=False)
+
+        # global_embeds = encoded.image_embeds
+        # global_embeds = global_embeds.unsqueeze(-2)
+        # prompt = ""
+
+        # ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        # encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+        # cak = dict(cond_lat=cond_lat)
+        image = torch.from_numpy(numpy.array(image)).to(self.vae.device) / 255.0
+        image = image.permute(2, 0, 1).unsqueeze(0)
+        if self.vae.dtype == torch.float16:
+            image = image.half()
+        # image = image.permute(2, 0, 1).to(self.vae.device).to(self.vae.dtype).unsqueeze(0)
+
+        latents = self.encode_target_images(image)
+        if cond_image is not None:
+            cond_image = to_rgb_image(cond_image)
+            cond_image = (
+                torch.from_numpy(numpy.array(cond_image)).to(self.vae.device) / 255.0
+            )
+            cond_image = cond_image.permute(2, 0, 1).unsqueeze(0)
+            if self.vae.dtype == torch.float16:
+                cond_image = cond_image.half()
+            cond_lat = self.encode_condition_image(cond_image)
+        else:
+            cond_lat = self.encode_condition_image(torch.zeros_like(image)).to(
+                self.vae.device
+            )
+        cak = dict(cond_lat=cond_lat, dont_forward_cond_state=True)
+        latents = self.forward_sdedit(
+            latents,
+            cross_attention_kwargs=cak,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=1,
+            prompt_embeds=encoder_hidden_states,
+            num_inference_steps=num_inference_steps,
+            output_type="latent",
+            width=width,
+            height=height,
+            edit_strength=edit_strength,
+            **kwargs,
+        ).images
+        # latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = unscale_image(
+                self.vae.decode(
+                    latents / self.vae.config.scaling_factor, return_dict=False
+                )[0]
+            )
+        else:
+            image = latents
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+    @torch.no_grad()
+    def refine(
+        self,
+        image: Image.Image = None,
+        edit_image: Image.Image = None,
+        prompt: Optional[str] = "",
+        *args,
+        output_type: Optional[str] = "pil",
+        width=640,
+        height=960,
+        num_inference_steps=28,
+        edit_strength=1.0,
+        return_dict=True,
+        guidance_scale=4.0,
+        **kwargs,
+    ):
+        self.prepare()
+        if image is None:
+            raise ValueError(
+                "Inputting embeddings not supported for this pipeline. Please pass an image."
+            )
+        assert not isinstance(image, torch.Tensor)
+        image = to_rgb_image(image)
+
+        # cond_lat = self.encode_condition_image(image_guidance)
+        if hasattr(self, "encode_prompt"):
+            encoder_hidden_states = self.encode_prompt(prompt, self.device, 1, False)[0]
+        else:
+            encoder_hidden_states = self._encode_prompt(prompt, self.device, 1, False)
+        # negative_lat = self.encode_condition_image(torch.zeros_like(image_guidance))
+        # cond_lat = torch.cat([negative_lat, cond_lat])
+        # encoded = self.vision_encoder(image_guidance_2, output_hidden_states=False)
+
+        # global_embeds = encoded.image_embeds
+        # global_embeds = global_embeds.unsqueeze(-2)
+        # prompt = ""
+
+        # ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        # encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+        # cak = dict(cond_lat=cond_lat)
+        latents_edit = None
+        if edit_image is not None:
+            edit_image = to_rgb_image(edit_image)
+            edit_image = (
+                torch.from_numpy(numpy.array(edit_image)).to(self.vae.device) / 255.0
+            )
+            edit_image = edit_image.permute(2, 0, 1).unsqueeze(0)
+            if self.vae.dtype == torch.float16:
+                edit_image = edit_image.half()
+            latents_edit = self.encode_target_images(edit_image)
+        image = torch.from_numpy(numpy.array(image)).to(self.vae.device) / 255.0
+        image = image.permute(2, 0, 1).unsqueeze(0)
+        if self.vae.dtype == torch.float16:
+            image = image.half()
+        # image = torch.nn.functional.interpolate(
+            # image, (height*4, width*4), mode="bilinear", align_corners=False)
+        # image = image[...,:320,:320]
+        height, width = image.shape[-2:]
+        # image = image[...,:640,:]
+        # image[...,:320,:] = torch.ones_like(image[...,:320,:])
+        # image = image.permute(2, 0, 1).to(self.vae.device).to(self.vae.dtype).unsqueeze(0)
+        # height = height * 4
+        # width = width * 4
+        latents = self.encode_target_images(image)
+        # latents[...,-40:,:] = torch.randn_like(latents[...,-40:,:])
+
+        cond_lat = self.encode_condition_image(torch.zeros_like(image)).to(
+            self.vae.device
+        )
+        cak = dict(cond_lat=cond_lat, dont_forward_cond_state=True)
+        latents = self.forward_pipeline(
+            latents_edit,
+            latents,
+            cross_attention_kwargs=cak,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=1,
+            prompt_embeds=encoder_hidden_states,
+            num_inference_steps=num_inference_steps,
+            output_type="latent",
+            width=width,
+            height=height,
+            edit_strength=edit_strength,
+            **kwargs,
+        ).images
+        # latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = unscale_image(
+                self.vae.decode(
+                    latents / self.vae.config.scaling_factor, return_dict=False
+                )[0]
+            )
+        else:
+            image = latents
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+        timestep=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=dtype
+            )
+            # scale the initial noise by the standard deviation required by the scheduler
+            latents = latents * self.scheduler.init_noise_sigma
+
+        else:
+            if timestep is None:
+                raise ValueError(
+                    "When passing `latents` you also need to pass `timestep`."
+                )
+            latents = latents.to(device)
+            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            # get latents
+            latents = self.scheduler.add_noise(latents, noise, timestep)
+
+        return latents
+
+    @torch.no_grad()
+    def forward_sdedit(
+        self,
+        latents: torch.Tensor,
+        cross_attention_kwargs: dict,
+        guidance_scale: float,
+        num_images_per_prompt: int,
+        prompt_embeds,
+        num_inference_steps: int,
+        output_type: str,
+        width: int,
+        height: int,
+        edit_strength: float = 1.0,
+    ):
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        batch_size = prompt_embeds.shape[0]
+        generator = torch.Generator(device=latents.device)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None)
+            if cross_attention_kwargs is not None
+            else None
+        )
+        prompt_embeds = self._encode_prompt(
+            None,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            None,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=None,
+            lora_scale=text_encoder_lora_scale,
+        )
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # self.scheduler.timesteps = self.scheduler.timesteps
+        timesteps = self.scheduler.timesteps
+        timesteps = reversed(reversed(timesteps)[: int(edit_strength * len(timesteps))])
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+            timesteps[0:1],
+        )
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, 0.0)
+        # if do_classifier_free_guidance:
+            # cond_latent = cond_latent.expand(batch_size * 2, -1, -1, -1)
+
+        # 7. Denoising loop
+        num_warmup_steps = 0
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                )
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+                # latent_model_input = 
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    return_dict=False,
+                )[0]
+                # exit(0)/
+                
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+
+                # 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()
+        latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = self.vae.decode(
+                latents / self.vae.config.scaling_factor, return_dict=False
+            )[0]
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(
+            image, output_type=output_type, do_denormalize=do_denormalize
+        )
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept
+        )
+
+    @torch.no_grad()
+    def forward_pipeline(
+        self,
+        latents: torch.Tensor,
+        cond_latent: torch.Tensor,
+        cross_attention_kwargs: dict,
+        guidance_scale: float,
+        num_images_per_prompt: int,
+        prompt_embeds,
+        num_inference_steps: int,
+        output_type: str,
+        width: int,
+        height: int,
+        edit_strength: float = 1.0,
+    ):
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        batch_size = 1
+        generator = torch.Generator(device=cond_latent.device)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None)
+            if cross_attention_kwargs is not None
+            else None
+        )
+        prompt_embeds = self._encode_prompt(
+            None,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            None,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=None,
+            lora_scale=text_encoder_lora_scale,
+        )
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # self.scheduler.timesteps = self.scheduler.timesteps
+        timesteps = self.scheduler.timesteps
+        timesteps = reversed(reversed(timesteps)[: int(edit_strength * len(timesteps))])
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels // 2
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+            timesteps[0:1],
+        )
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, 0.0)
+        if do_classifier_free_guidance:
+            cond_latent = cond_latent.expand(batch_size * 2, -1, -1, -1)
+
+        # 7. Denoising loop
+        num_warmup_steps = 0
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                )
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+                latent_model_input = torch.cat([latent_model_input, cond_latent], dim=1)
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+
+                # 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()
+        latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = self.vae.decode(
+                latents / self.vae.config.scaling_factor, return_dict=False
+            )[0]
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(
+            image, output_type=output_type, do_denormalize=do_denormalize
+        )
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept
+        )
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Image.Image = None,
+        source_image: Image.Image = None,
+        prompt="",
+        *args,
+        num_images_per_prompt: Optional[int] = 1,
+        guidance_scale=4.0,
+        depth_image: Image.Image = None,
+        output_type: Optional[str] = "pil",
+        width=640,
+        height=960,
+        num_inference_steps=28,
+        return_dict=True,
+        **kwargs,
+    ):
+        self.prepare()
+        if image is None:
+            raise ValueError(
+                "Inputting embeddings not supported for this pipeline. Please pass an image."
+            )
+        assert not isinstance(image, torch.Tensor)
+        image = to_rgb_image(image)
+        image_1 = self.feature_extractor_vae(
+            images=image, return_tensors="pt"
+        ).pixel_values
+        image_2 = self.feature_extractor_clip(
+            images=image, return_tensors="pt"
+        ).pixel_values
+        # image_source = to_rgb_image(source_image)
+        # image_source_latents = self.feature_extractor_vae(images=image_source, return_tensors="pt")
+        if depth_image is not None and hasattr(self.unet, "controlnet"):
+            depth_image = to_rgb_image(depth_image)
+            depth_image = self.depth_transforms_multi(depth_image).to(
+                device=self.unet.controlnet.device, dtype=self.unet.controlnet.dtype
+            )
+        image = image_1.to(device=self.vae.device, dtype=self.vae.dtype)
+        image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)
+        cond_lat = self.encode_condition_image(image)
+        if guidance_scale > 1:
+            negative_lat = self.encode_condition_image(torch.zeros_like(image))
+            cond_lat = torch.cat([negative_lat, cond_lat])
+        encoded = self.vision_encoder(image_2, output_hidden_states=False)
+        global_embeds = encoded.image_embeds
+        global_embeds = global_embeds.unsqueeze(-2)
+
+        if hasattr(self, "encode_prompt"):
+            encoder_hidden_states = self.encode_prompt(
+                prompt, self.device, num_images_per_prompt, False
+            )[0]
+        else:
+            encoder_hidden_states = self._encode_prompt(
+                prompt, self.device, num_images_per_prompt, False
+            )
+        ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
+        cak = dict(cond_lat=cond_lat)
+        if hasattr(self.unet, "controlnet"):
+            cak["control_depth"] = depth_image
+        latents: torch.Tensor = (
+            super()
+            .__call__(
+                None,
+                *args,
+                cross_attention_kwargs=cak,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=num_images_per_prompt,
+                prompt_embeds=encoder_hidden_states,
+                num_inference_steps=num_inference_steps,
+                output_type="latent",
+                width=width,
+                height=height,
+                latents=None,
+                **kwargs,
+            )
+            .images
+        )
+        latents = unscale_latents(latents)
+        if not output_type == "latent":
+            image = unscale_image(
+                self.vae.decode(
+                    latents / self.vae.config.scaling_factor, return_dict=False
+                )[0]
+            )
+        else:
+            image = latents
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)