chores: rebase commits

threestudio/models/guidance/__init__.py

@@ -0,0 +1,13 @@
+from . import (
+    controlnet_guidance,
+    controlnet_reg_guidance,
+    deep_floyd_guidance,
+    stable_diffusion_guidance,
+    stable_diffusion_unified_guidance,
+    stable_diffusion_vsd_guidance,
+    stable_diffusion_bsd_guidance,
+    stable_zero123_guidance,
+    zero123_guidance,
+    zero123_unified_guidance,
+    clip_guidance,
+)

threestudio/models/guidance/clip_guidance.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+import clip
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.typing import *
+
+
+@threestudio.register("clip-guidance")
+class CLIPGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        pretrained_model_name_or_path: str = "ViT-B/16"
+        view_dependent_prompting: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading CLIP ...")
+        self.clip_model, self.clip_preprocess = clip.load(
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            jit=False,
+            download_root=self.cfg.cache_dir
+        )
+
+        self.aug = T.Compose([
+            T.Resize((224, 224)),
+            T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+        threestudio.info(f"Loaded CLIP!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def get_embedding(self, input_value, is_text=True):
+        if is_text:
+            value = clip.tokenize(input_value).to(self.device)
+            z = self.clip_model.encode_text(value)
+        else:
+            input_value = self.aug(input_value)
+            z = self.clip_model.encode_image(input_value)
+
+        return z / z.norm(dim=-1, keepdim=True)
+
+    def get_loss(self, image_z, clip_z, loss_type='similarity_score', use_mean=True):
+        if loss_type == 'similarity_score':
+            loss = -((image_z * clip_z).sum(-1))
+        elif loss_type == 'spherical_dist':
+            image_z, clip_z = F.normalize(image_z, dim=-1), F.normalize(clip_z, dim=-1)
+            loss = ((image_z - clip_z).norm(dim=-1).div(2).arcsin().pow(2).mul(2))
+        else:
+            raise NotImplementedError
+
+        return loss.mean() if use_mean else loss
+
+    def __call__(
+        self,
+        pred_rgb: Float[Tensor, "B H W C"],
+        gt_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        embedding_type: str = 'both',
+        loss_type: Optional[str] = 'similarity_score',
+        **kwargs,
+    ):
+        clip_text_loss, clip_img_loss = 0, 0
+
+        if embedding_type in ('both', 'text'):
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            ).chunk(2)[0]
+            clip_text_loss = self.get_loss(self.get_embedding(pred_rgb, is_text=False), text_embeddings, loss_type=loss_type)
+
+        if embedding_type in ('both', 'img'):
+            clip_img_loss = self.get_loss(self.get_embedding(pred_rgb, is_text=False), self.get_embedding(gt_rgb, is_text=False), loss_type=loss_type)
+
+        return clip_text_loss + clip_img_loss

threestudio/models/guidance/controlnet_guidance.py

@@ -0,0 +1,517 @@
+import os
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from controlnet_aux import CannyDetector, NormalBaeDetector
+from diffusers import ControlNetModel, DDIMScheduler, StableDiffusionControlNetPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.perceptual import PerceptualLoss
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-controlnet-guidance")
+class ControlNetGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        pretrained_model_name_or_path: str = "SG161222/Realistic_Vision_V2.0"
+        ddim_scheduler_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        control_type: str = "normal"  # normal/canny
+
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 7.5
+        condition_scale: float = 1.5
+        grad_clip: Optional[Any] = None
+        half_precision_weights: bool = True
+
+        fixed_size: int = -1
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        diffusion_steps: int = 20
+
+        use_sds: bool = False
+
+        use_du: bool = False
+        per_du_step: int = 10
+        start_du_step: int = 1000
+        cache_du: bool = False
+
+        # Canny threshold
+        canny_lower_bound: int = 50
+        canny_upper_bound: int = 100
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading ControlNet ...")
+
+        controlnet_name_or_path: str
+        if self.cfg.control_type in ("normal", "input_normal"):
+            controlnet_name_or_path = "lllyasviel/control_v11p_sd15_normalbae"
+        elif self.cfg.control_type == "canny":
+            controlnet_name_or_path = "lllyasviel/control_v11p_sd15_canny"
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        pipe_kwargs = {
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+        }
+
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_name_or_path,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+        )
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, controlnet=controlnet, **pipe_kwargs
+        ).to(self.device)
+        self.scheduler = DDIMScheduler.from_pretrained(
+            self.cfg.ddim_scheduler_name_or_path,
+            subfolder="scheduler",
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+        )
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        # Create model
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+        self.controlnet = self.pipe.controlnet.eval()
+
+        if self.cfg.control_type == "normal":
+            self.preprocessor = NormalBaeDetector.from_pretrained(
+                "lllyasviel/Annotators"
+            )
+            self.preprocessor.model.to(self.device)
+        elif self.cfg.control_type == "canny":
+            self.preprocessor = CannyDetector()
+
+        for p in self.vae.parameters():
+            p.requires_grad_(False)
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        if self.cfg.use_du:
+            if self.cfg.cache_du:
+                self.edit_frames = {}
+            self.perceptual_loss = PerceptualLoss().eval().to(self.device)
+
+        threestudio.info(f"Loaded ControlNet!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_controlnet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        image_cond: Float[Tensor, "..."],
+        condition_scale: float,
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        return self.controlnet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            controlnet_cond=image_cond.to(self.weights_dtype),
+            conditioning_scale=condition_scale,
+            return_dict=False,
+        )
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_control_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        cross_attention_kwargs,
+        down_block_additional_residuals,
+        mid_block_additional_residual,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_cond_images(
+        self, imgs: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.mode()
+        uncond_image_latents = torch.zeros_like(latents)
+        latents = torch.cat([latents, latents, uncond_image_latents], dim=0)
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self, latents: Float[Tensor, "B 4 DH DW"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        input_dtype = latents.dtype
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def edit_latents(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 DH DW"],
+        image_cond: Float[Tensor, "B 3 H W"],
+        t: Int[Tensor, "B"],
+        mask = None
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        self.scheduler.config.num_train_timesteps = t.item()
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+        if mask is not None:
+            mask = F.interpolate(mask, (latents.shape[-2], latents.shape[-1]), mode='bilinear')
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents = self.scheduler.add_noise(latents, noise, t)  # type: ignore
+
+            # sections of code used from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
+            threestudio.debug("Start editing...")
+            for i, t in enumerate(self.scheduler.timesteps):
+                # predict the noise residual with unet, NO grad!
+                with torch.no_grad():
+                    # pred noise
+                    latent_model_input = torch.cat([latents] * 2)
+                    (
+                        down_block_res_samples,
+                        mid_block_res_sample,
+                    ) = self.forward_controlnet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        image_cond=image_cond,
+                        condition_scale=self.cfg.condition_scale,
+                    )
+
+                    noise_pred = self.forward_control_unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs=None,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                    )
+                # perform classifier-free guidance
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+                if mask is not None:
+                    noise_pred = mask * noise_pred + (1 - mask) * noise
+                # get previous sample, continue loop
+                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+            threestudio.debug("Editing finished.")
+        return latents
+
+    def prepare_image_cond(self, cond_rgb: Float[Tensor, "B H W C"]):
+        if self.cfg.control_type == "normal":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            detected_map = self.preprocessor(cond_rgb)
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "canny":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            blurred_img = cv2.blur(cond_rgb, ksize=(5, 5))
+            detected_map = self.preprocessor(
+                blurred_img, self.cfg.canny_lower_bound, self.cfg.canny_upper_bound
+            )
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            # control = control.unsqueeze(-1).repeat(1, 1, 3)
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "input_normal":
+            cond_rgb[..., 0] = (
+                1 - cond_rgb[..., 0]
+            )  # Flip the sign on the x-axis to match bae system
+            control = cond_rgb.permute(0, 3, 1, 2)
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+
+        return control
+
+    def compute_grad_sds(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 DH DW"],
+        image_cond: Float[Tensor, "B 3 H W"],
+        t: Int[Tensor, "B"],
+    ):
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            down_block_res_samples, mid_block_res_sample = self.forward_controlnet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                image_cond=image_cond,
+                condition_scale=self.cfg.condition_scale,
+            )
+
+            noise_pred = self.forward_control_unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                cross_attention_kwargs=None,
+                down_block_additional_residuals=down_block_res_samples,
+                mid_block_additional_residual=mid_block_res_sample,
+            )
+
+        # perform classifier-free guidance
+        noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        return grad
+
+    def compute_grad_du(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        rgb_BCHW_HW8: Float[Tensor, "B 3 RH RW"],
+        cond_feature: Float[Tensor, "B 3 RH RW"],
+        cond_rgb: Float[Tensor, "B H W 3"],
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        mask = None,
+        **kwargs,
+    ):
+        batch_size, _, RH, RW = cond_feature.shape
+        assert batch_size == 1
+
+        origin_gt_rgb = F.interpolate(
+            cond_rgb.permute(0, 3, 1, 2), (RH, RW), mode="bilinear"
+        ).permute(0, 2, 3, 1)
+        need_diffusion = (
+            self.global_step % self.cfg.per_du_step == 0
+            and self.global_step > self.cfg.start_du_step
+        )
+        if self.cfg.cache_du:
+            if torch.is_tensor(kwargs["index"]):
+                batch_index = kwargs["index"].item()
+            else:
+                batch_index = kwargs["index"]
+            if (
+                not (batch_index in self.edit_frames)
+            ) and self.global_step > self.cfg.start_du_step:
+                need_diffusion = True
+        need_loss = self.cfg.cache_du or need_diffusion
+        guidance_out = {}
+
+        if need_diffusion:
+            t = torch.randint(
+                self.min_step,
+                self.max_step,
+                [1],
+                dtype=torch.long,
+                device=self.device,
+            )
+            print("t:", t)
+            edit_latents = self.edit_latents(text_embeddings, latents, cond_feature, t, mask)
+            edit_images = self.decode_latents(edit_latents)
+            edit_images = F.interpolate(
+                edit_images, (RH, RW), mode="bilinear"
+            ).permute(0, 2, 3, 1)
+            self.edit_images = edit_images
+            if self.cfg.cache_du:
+                self.edit_frames[batch_index] = edit_images.detach().cpu()
+
+        if need_loss:
+            if self.cfg.cache_du:
+                if batch_index in self.edit_frames:
+                    gt_rgb = self.edit_frames[batch_index].to(cond_feature.device)
+                else:
+                    gt_rgb = origin_gt_rgb
+            else:
+                gt_rgb = edit_images
+
+            import cv2
+            import numpy as np
+
+            temp = (edit_images.detach().cpu()[0].numpy() * 255).astype(np.uint8)
+            cv2.imwrite(".threestudio_cache/test.jpg", temp[:, :, ::-1])
+            
+            guidance_out.update(
+                {
+                    "loss_l1": torch.nn.functional.l1_loss(
+                        rgb_BCHW_HW8, gt_rgb.permute(0, 3, 1, 2), reduction="sum"
+                    ),
+                    "loss_p": self.perceptual_loss(
+                        rgb_BCHW_HW8.contiguous(),
+                        gt_rgb.permute(0, 3, 1, 2).contiguous(),
+                    ).sum(),
+                }
+            )
+
+        return guidance_out
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        cond_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        mask = None,
+        **kwargs,
+    ):
+        batch_size, H, W, _ = rgb.shape
+        assert batch_size == 1
+        assert rgb.shape[:-1] == cond_rgb.shape[:-1]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        if mask is not None: mask = mask.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 DH DW"]
+        if self.cfg.fixed_size > 0:
+            RH, RW = self.cfg.fixed_size, self.cfg.fixed_size
+        else:
+            RH, RW = H // 8 * 8, W // 8 * 8
+        rgb_BCHW_HW8 = F.interpolate(
+            rgb_BCHW, (RH, RW), mode="bilinear", align_corners=False
+        )
+        latents = self.encode_images(rgb_BCHW_HW8)
+
+        image_cond = self.prepare_image_cond(cond_rgb)
+        image_cond = F.interpolate(
+            image_cond, (RH, RW), mode="bilinear", align_corners=False
+        )
+
+        temp = torch.zeros(1).to(rgb.device)
+        azimuth = kwargs.get("azimuth", temp)
+        camera_distance = kwargs.get("camera_distance", temp)
+        view_dependent_prompt = kwargs.get("view_dependent_prompt", False)
+        text_embeddings = prompt_utils.get_text_embeddings(temp, azimuth, camera_distance, view_dependent_prompt) # FIXME: change to view-conditioned prompt
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+
+        guidance_out = {}
+        if self.cfg.use_sds:
+            grad = self.compute_grad_sds(text_embeddings, latents, image_cond, t)
+            grad = torch.nan_to_num(grad)
+            if self.grad_clip_val is not None:
+                grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+            target = (latents - grad).detach()
+            loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+            guidance_out.update(
+                {
+                    "loss_sds": loss_sds,
+                    "grad_norm": grad.norm(),
+                    "min_step": self.min_step,
+                    "max_step": self.max_step,
+                }
+            )
+
+        if self.cfg.use_du:
+            grad = self.compute_grad_du(
+                latents, rgb_BCHW_HW8, image_cond, cond_rgb, text_embeddings, mask, **kwargs
+            )
+            guidance_out.update(grad)
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+        self.global_step = global_step

threestudio/models/guidance/controlnet_reg_guidance.py

@@ -0,0 +1,454 @@
+import os
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from controlnet_aux import CannyDetector, NormalBaeDetector
+from diffusers import ControlNetModel, DDIMScheduler, StableDiffusionControlNetPipeline, DPMSolverMultistepScheduler
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-controlnet-reg-guidance")
+class ControlNetGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "SG161222/Realistic_Vision_V2.0"
+        ddim_scheduler_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        control_type: str = "normal"  # normal/canny
+
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 7.5
+        condition_scale: float = 1.5
+        grad_clip: Optional[Any] = None
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        diffusion_steps: int = 20
+
+        use_sds: bool = False
+
+        # Canny threshold
+        canny_lower_bound: int = 50
+        canny_upper_bound: int = 100
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading ControlNet ...")
+
+        self.weights_dtype = torch.float16 if self.cfg.half_precision_weights else torch.float32
+
+        self.preprocessor, controlnet_name_or_path = self.get_preprocessor_and_controlnet()
+        
+        pipe_kwargs = self.configure_pipeline()
+
+        self.load_models(pipe_kwargs, controlnet_name_or_path)
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.scheduler = self.pipe.scheduler
+
+        self.check_memory_efficiency_conditions()
+
+        self.set_min_max_steps()
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device)
+        self.grad_clip_val = None
+
+        threestudio.info(f"Loaded ControlNet!")
+
+    def get_preprocessor_and_controlnet(self):
+        if self.cfg.control_type in ("normal", "input_normal"):
+            if self.cfg.pretrained_model_name_or_path == "SG161222/Realistic_Vision_V2.0":
+                controlnet_name_or_path = "lllyasviel/control_v11p_sd15_normalbae"
+            else:
+                controlnet_name_or_path = "thibaud/controlnet-sd21-normalbae-diffusers"
+            preprocessor = NormalBaeDetector.from_pretrained("lllyasviel/Annotators", cache_dir=self.cfg.cache_dir)
+            preprocessor.model.to(self.device)
+        elif self.cfg.control_type == "canny" or self.cfg.control_type == "canny2":
+            controlnet_name_or_path = self.get_canny_controlnet()
+            preprocessor = CannyDetector()
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+        return preprocessor, controlnet_name_or_path
+
+    def get_canny_controlnet(self):
+        if self.cfg.control_type == "canny":
+            return "lllyasviel/control_v11p_sd15_canny"
+        elif self.cfg.control_type == "canny2":
+            return "thepowefuldeez/sd21-controlnet-canny"
+
+    def configure_pipeline(self):
+        return {
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only
+        }
+
+    def load_models(self, pipe_kwargs, controlnet_name_or_path):
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_name_or_path,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        )
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, controlnet=controlnet, **pipe_kwargs
+        ).to(self.device)
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            self.cfg.ddim_scheduler_name_or_path,
+            subfolder="scheduler",
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        )
+
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+        self.controlnet = self.pipe.controlnet.eval()
+
+    def check_memory_efficiency_conditions(self):
+        if self.cfg.enable_memory_efficient_attention:
+            self.memory_efficiency_status()
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+    def memory_efficiency_status(self):
+        if parse_version(torch.__version__) >= parse_version("2"):
+            threestudio.info("PyTorch2.0 uses memory efficient attention by default.")
+        elif not is_xformers_available():
+            threestudio.warn("xformers is not available, memory efficient attention is not enabled.")
+        else:
+            self.pipe.enable_xformers_memory_efficient_attention()
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_controlnet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        image_cond: Float[Tensor, "..."],
+        condition_scale: float,
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        return self.controlnet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            controlnet_cond=image_cond.to(self.weights_dtype),
+            conditioning_scale=condition_scale,
+            return_dict=False,
+        )
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_control_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        cross_attention_kwargs,
+        down_block_additional_residuals,
+        mid_block_additional_residual,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_cond_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.mode()
+        uncond_image_latents = torch.zeros_like(latents)
+        latents = torch.cat([latents, latents, uncond_image_latents], dim=0)
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        latent_height: int = 64,
+        latent_width: int = 64,
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        latents = F.interpolate(
+            latents, (latent_height, latent_width), mode="bilinear", align_corners=False
+        )
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def edit_latents(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 64 64"],
+        image_cond: Float[Tensor, "B 3 512 512"],
+        t: Int[Tensor, "B"],
+        mask=None
+    ) -> Float[Tensor, "B 4 64 64"]:
+        batch_size = t.shape[0]
+        self.scheduler.set_timesteps(num_inference_steps=self.cfg.diffusion_steps)
+        init_timestep = max(1, min(int(self.cfg.diffusion_steps * t[0].item() / self.num_train_timesteps), self.cfg.diffusion_steps))
+        t_start = max(self.cfg.diffusion_steps - init_timestep, 0)
+        latent_timestep = self.scheduler.timesteps[t_start : t_start + 1].repeat(batch_size)
+        B, _, DH, DW = latents.shape
+        origin_latents = latents.clone()
+        if mask is not None:
+            mask = F.interpolate(mask, (DH, DW), mode="bilinear", antialias=True)
+        
+        with torch.no_grad():
+            # sections of code used from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
+            noise = torch.randn_like(latents)
+            latents = self.scheduler.add_noise(latents, noise, latent_timestep)  # type: ignore
+            threestudio.debug("Start editing...")
+            for i, step in enumerate(range(t_start, self.cfg.diffusion_steps)):
+                timestep = self.scheduler.timesteps[step]
+                # predict the noise residual with unet, NO grad!
+                with torch.no_grad():
+                    # pred noise
+                    latent_model_input = torch.cat([latents] * 2)
+                    (
+                        down_block_res_samples,
+                        mid_block_res_sample,
+                    ) = self.forward_controlnet(
+                        latent_model_input,
+                        timestep,
+                        encoder_hidden_states=text_embeddings,
+                        image_cond=image_cond,
+                        condition_scale=self.cfg.condition_scale,
+                    )
+
+                    noise_pred = self.forward_control_unet(
+                        latent_model_input,
+                        timestep,
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs=None,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                    )
+                # perform classifier-free guidance
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+                if mask is not None:
+                    noise_pred = noise_pred * mask + (1-mask) * noise
+                
+                latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample
+            threestudio.debug("Editing finished.")
+        return latents
+
+    def prepare_image_cond(self, cond_rgb: Float[Tensor, "B H W C"]):
+        if self.cfg.control_type == "normal":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            detected_map = self.preprocessor(cond_rgb)
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "canny" or self.cfg.control_type == "canny2":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            blurred_img = cv2.blur(cond_rgb, ksize=(5, 5))
+            detected_map = self.preprocessor(
+                blurred_img, self.cfg.canny_lower_bound, self.cfg.canny_upper_bound
+            )
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(-1).repeat(1, 1, 3)
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "input_normal":
+            cond_rgb[..., 0] = (
+                1 - cond_rgb[..., 0]
+            )  # Flip the sign on the x-axis to match bae system
+            control = cond_rgb.permute(0, 3, 1, 2)
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+
+        return F.interpolate(control, (512, 512), mode="bilinear", align_corners=False)
+
+    def compute_grad_sds(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 64 64"],
+        image_cond: Float[Tensor, "B 3 512 512"],
+        t: Int[Tensor, "B"],
+    ):
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            down_block_res_samples, mid_block_res_sample = self.forward_controlnet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                image_cond=image_cond,
+                condition_scale=self.cfg.condition_scale,
+            )
+
+            noise_pred = self.forward_control_unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                cross_attention_kwargs=None,
+                down_block_additional_residuals=down_block_res_samples,
+                mid_block_additional_residual=mid_block_res_sample,
+            )
+
+        # perform classifier-free guidance
+        noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        return grad
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        cond_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        mask: Float[Tensor, "B H W C"],
+        **kwargs,
+    ):
+        
+        batch_size, H, W, _ = rgb.shape
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        rgb_BCHW_512 = F.interpolate(
+            rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+        )
+        latents = self.encode_images(rgb_BCHW_512)
+
+        image_cond = self.prepare_image_cond(cond_rgb)
+
+        temp = torch.zeros(1).to(rgb.device)
+        text_embeddings = prompt_utils.get_text_embeddings(temp, temp, temp, False)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        if self.cfg.use_sds:
+            grad = self.compute_grad_sds(text_embeddings, latents, image_cond, t)
+            grad = torch.nan_to_num(grad)
+            if self.grad_clip_val is not None:
+                grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+            target = (latents - grad).detach()
+            loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+            return {
+                "loss_sds": loss_sds,
+                "grad_norm": grad.norm(),
+                "min_step": self.min_step,
+                "max_step": self.max_step,
+            }
+        else:
+            
+            if mask is not None: mask = mask.permute(0, 3, 1, 2)
+            edit_latents = self.edit_latents(text_embeddings, latents, image_cond, t, mask)
+            edit_images = self.decode_latents(edit_latents)
+            edit_images = F.interpolate(edit_images, (H, W), mode="bilinear")
+
+            return {"edit_images": edit_images.permute(0, 2, 3, 1),
+                    "edit_latents": edit_latents}
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+
+if __name__ == "__main__":
+    from threestudio.utils.config import ExperimentConfig, load_config
+    from threestudio.utils.typing import Optional
+
+    cfg = load_config("configs/experimental/controlnet-normal.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(
+        cfg.system.prompt_processor
+    )
+
+    rgb_image = cv2.imread("assets/face.jpg")[:, :, ::-1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (512, 512))
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device)
+    prompt_utils = prompt_processor()
+    guidance_out = guidance(rgb_image, rgb_image, prompt_utils)
+    edit_image = (
+        (guidance_out["edit_images"][0].detach().cpu().clip(0, 1).numpy() * 255)
+        .astype(np.uint8)[:, :, ::-1]
+        .copy()
+    )
+    os.makedirs(".threestudio_cache", exist_ok=True)
+    cv2.imwrite(".threestudio_cache/edit_image.jpg", edit_image)

threestudio/models/guidance/deep_floyd_guidance.py

@@ -0,0 +1,582 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import IFPipeline, DDPMScheduler
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("deep-floyd-guidance")
+class DeepFloydGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "DeepFloyd/IF-I-XL-v1.0"
+        # FIXME: xformers error
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = True
+        guidance_scale: float = 20.0
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        time_prior: Optional[Any] = None # [w1,w2,s1,s2]
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        weighting_strategy: str = "sds"
+
+        view_dependent_prompting: bool = True
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+        lora_weights_path: Optional[str] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Deep Floyd ...")
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        # Create model
+        self.pipe = IFPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            text_encoder=None,
+            safety_checker=None,
+            watermarker=None,
+            feature_extractor=None,
+            requires_safety_checker=False,
+            variant="fp16" if self.cfg.half_precision_weights else None,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        ).to(self.device)
+
+        # Load lora weights
+        if self.cfg.lora_weights_path is not None:
+            self.pipe.load_lora_weights(self.cfg.lora_weights_path)
+            self.pipe.scheduler = self.pipe.scheduler.__class__.from_config(self.pipe.scheduler.config, variance_type="fixed_small")
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                threestudio.warn(
+                    f"Use DeepFloyd with xformers may raise error, see https://github.com/deep-floyd/IF/issues/52 to track this problem."
+                )
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        self.unet = self.pipe.unet.eval()
+
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        self.scheduler = self.pipe.scheduler
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+        if self.cfg.time_prior is not None:
+            m1, m2, s1, s2 = self.cfg.time_prior
+            weights = torch.cat(
+                (
+                    torch.exp(
+                        -((torch.arange(self.num_train_timesteps, m1, -1) - m1) ** 2)
+                        / (2 * s1**2)
+                    ),
+                    torch.ones(m1 - m2 + 1),
+                    torch.exp(
+                        -((torch.arange(m2 - 1, 0, -1) - m2) ** 2) / (2 * s2**2)
+                    ),
+                )
+            )
+            weights = weights / torch.sum(weights)
+            self.time_prior_acc_weights = torch.cumsum(weights, dim=0)
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        threestudio.info(f"Loaded Deep Floyd!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+        ).sample.to(input_dtype)
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        current_step_ratio=None,
+        mask: Float[Tensor, "B H W 1"] = None,
+        rgb_as_latents=False,
+        guidance_eval=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+
+        if mask is not None: 
+            mask = mask.permute(0, 3, 1, 2)
+            mask = F.interpolate(
+                mask, (64, 64), mode="bilinear", align_corners=False
+            )
+
+        assert rgb_as_latents == False, f"No latent space in {self.__class__.__name__}"
+        rgb_BCHW = rgb_BCHW * 2.0 - 1.0  # scale to [-1, 1] to match the diffusion range
+        latents = F.interpolate(
+            rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+        )
+
+        if self.cfg.time_prior is not None:
+            time_index = torch.where(
+                (self.time_prior_acc_weights - current_step_ratio) > 0
+            )[0][0]
+            if time_index == 0 or torch.abs(
+                self.time_prior_acc_weights[time_index] - current_step_ratio
+            ) < torch.abs(
+                self.time_prior_acc_weights[time_index - 1] - current_step_ratio
+            ):
+                t = self.num_train_timesteps - time_index
+            else:
+                t = self.num_train_timesteps - time_index + 1
+            t = torch.clip(t, self.min_step, self.max_step + 1)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+
+        else:
+            # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+            t = torch.randint(
+                self.min_step,
+                self.max_step + 1,
+                [batch_size],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                if mask is not None:
+                    latents_noisy = (1 - mask) * latents + mask * latents_noisy
+                latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 6, 64, 64)
+
+            noise_pred_text, _ = noise_pred[:batch_size].split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred[batch_size : batch_size * 2].split(
+                3, dim=1
+            )
+            noise_pred_neg, _ = noise_pred[batch_size * 2 :].split(3, dim=1)
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                if mask is not None:
+                    latents_noisy = (1 - mask) * latents + mask * latents_noisy
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )  # (2B, 6, 64, 64)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        """
+        # thresholding, experimental
+        if self.cfg.thresholding:
+            assert batch_size == 1
+            noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)
+            noise_pred = custom_ddpm_step(self.scheduler,
+                noise_pred, int(t.item()), latents_noisy, **self.pipe.prepare_extra_step_kwargs(None, 0.0)
+            )
+        """
+
+        if self.cfg.weighting_strategy == "sds":
+            # w(t), sigma_t^2
+            w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        # # FIXME: Visualize inpainting results
+        # self.scheduler.set_timesteps(20)
+        # latents = latents_noisy
+        # for t in tqdm(self.scheduler.timesteps):
+        #         # pred noise
+        #         noise_pred = self.get_noise_pred(
+        #             latents, t, text_embeddings, prompt_utils.use_perp_neg, None
+        #         )
+        #         # get prev latent
+        #         prev_latents = latents
+        #         latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
+        #         if mask is not None:
+        #             latents = (1 - mask) * prev_latents + mask * latents
+
+        # denoised_img = (latents / 2 + 0.5).permute(0, 2, 3, 1)
+        # guidance_out.update(
+        #     {"denoised_img": denoised_img}
+        # )
+
+        if guidance_eval:
+            guidance_eval_utils = {
+                "use_perp_neg": prompt_utils.use_perp_neg,
+                "neg_guidance_weights": neg_guidance_weights,
+                "text_embeddings": text_embeddings,
+                "t_orig": t,
+                "latents_noisy": latents_noisy,
+                "noise_pred": torch.cat([noise_pred, predicted_variance], dim=1),
+            }
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_noise_pred(
+        self,
+        latents_noisy,
+        t,
+        text_embeddings,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        batch_size = latents_noisy.shape[0]
+        if use_perp_neg:
+            latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 4).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (4B, 6, 64, 64)
+
+            noise_pred_text, _ = noise_pred[:batch_size].split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred[batch_size : batch_size * 2].split(
+                3, dim=1
+            )
+            noise_pred_neg, _ = noise_pred[batch_size * 2 :].split(3, dim=1)
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 2).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (2B, 6, 64, 64)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return torch.cat([noise_pred, predicted_variance], dim=1)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(
+        self,
+        t_orig,
+        text_embeddings,
+        latents_noisy,
+        noise_pred,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = (latents_noisy[:bs] / 2 + 0.5).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1]
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = (latents_1step / 2 + 0.5).permute(0, 2, 3, 1)
+        imgs_1orig = (pred_1orig / 2 + 0.5).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            text_emb = (
+                text_embeddings[
+                    [b, b + len(idxs), b + 2 * len(idxs), b + 3 * len(idxs)], ...
+                ]
+                if use_perp_neg
+                else text_embeddings[[b, b + len(idxs)], ...]
+            )
+            neg_guid = neg_guidance_weights[b : b + 1] if use_perp_neg else None
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                noise_pred = self.get_noise_pred(
+                    latents, t, text_emb, use_perp_neg, neg_guid
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = (latents_final / 2 + 0.5).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+
+"""
+# used by thresholding, experimental
+def custom_ddpm_step(ddpm, model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, generator=None, return_dict: bool = True):
+    self = ddpm
+    t = timestep
+
+    prev_t = self.previous_timestep(t)
+
+    if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
+        model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
+    else:
+        predicted_variance = None
+
+    # 1. compute alphas, betas
+    alpha_prod_t = self.alphas_cumprod[t].item()
+    alpha_prod_t_prev = self.alphas_cumprod[prev_t].item() if prev_t >= 0 else 1.0
+    beta_prod_t = 1 - alpha_prod_t
+    beta_prod_t_prev = 1 - alpha_prod_t_prev
+    current_alpha_t = alpha_prod_t / alpha_prod_t_prev
+    current_beta_t = 1 - current_alpha_t
+
+    # 2. compute predicted original sample from predicted noise also called
+    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+    if self.config.prediction_type == "epsilon":
+        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+    elif self.config.prediction_type == "sample":
+        pred_original_sample = model_output
+    elif self.config.prediction_type == "v_prediction":
+        pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+    else:
+        raise ValueError(
+            f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+            " `v_prediction`  for the DDPMScheduler."
+        )
+
+    # 3. Clip or threshold "predicted x_0"
+    if self.config.thresholding:
+        pred_original_sample = self._threshold_sample(pred_original_sample)
+    elif self.config.clip_sample:
+        pred_original_sample = pred_original_sample.clamp(
+            -self.config.clip_sample_range, self.config.clip_sample_range
+        )
+
+    noise_thresholded = (sample - (alpha_prod_t ** 0.5) * pred_original_sample) / (beta_prod_t ** 0.5)
+    return noise_thresholded
+"""
+
+
+if __name__ == '__main__':
+    from threestudio.utils.config import load_config
+    import pytorch_lightning as pl
+    import numpy as np
+    import os
+    import cv2
+    cfg = load_config("configs/debugging/deepfloyd.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(cfg.system.prompt_processor)
+    prompt_utils = prompt_processor()
+    temp = torch.zeros(1).to(guidance.device)
+    # rgb_image = guidance.sample(prompt_utils, temp, temp, temp, seed=cfg.seed)
+    # rgb_image = (rgb_image[0].detach().cpu().clip(0, 1).numpy()*255).astype(np.uint8)[:, :, ::-1].copy()
+    # os.makedirs('.threestudio_cache', exist_ok=True)
+    # cv2.imwrite('.threestudio_cache/diffusion_image.jpg', rgb_image)
+
+    ### inpaint
+    rgb_image = cv2.imread("assets/test.jpg")[:, :, ::-1].copy() / 255
+    mask_image = cv2.imread("assets/mask.png")[:, :, :1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (512, 512))
+    mask_image = cv2.resize(mask_image, (512, 512)).reshape(512, 512, 1)
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device)
+    mask_image = torch.FloatTensor(mask_image).unsqueeze(0).to(guidance.device)
+
+    guidance_out = guidance(rgb_image, prompt_utils, temp, temp, temp, mask=mask_image)
+    edit_image = (
+        (guidance_out["denoised_img"][0].detach().cpu().clip(0, 1).numpy() * 255)
+        .astype(np.uint8)[:, :, ::-1]
+        .copy()
+    )
+    os.makedirs(".threestudio_cache", exist_ok=True)
+    cv2.imwrite(".threestudio_cache/edit_image.jpg", edit_image)

threestudio/models/guidance/stable_diffusion_guidance.py

@@ -0,0 +1,632 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, cleanup, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-guidance")
+class StableDiffusionGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 100.0
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        time_prior: Optional[Any] = None # [w1,w2,s1,s2]
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+        max_step_percent_annealed: float = 0.5
+        anneal_start_step: Optional[int] = None
+
+        use_sjc: bool = False
+        var_red: bool = True
+        weighting_strategy: str = "sds"
+
+        token_merging: bool = False
+        token_merging_params: Optional[dict] = field(default_factory=dict)
+
+        view_dependent_prompting: bool = True
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Stable Diffusion ...")
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        pipe_kwargs = {
+            "tokenizer": None,
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only
+        }
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        del self.pipe.text_encoder
+        cleanup()
+
+        # Create model
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+
+        for p in self.vae.parameters():
+            p.requires_grad_(False)
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        if self.cfg.token_merging:
+            import tomesd
+
+            tomesd.apply_patch(self.unet, **self.cfg.token_merging_params)
+
+        if self.cfg.use_sjc:
+            # score jacobian chaining use DDPM
+            self.scheduler = DDPMScheduler.from_pretrained(
+                self.cfg.pretrained_model_name_or_path,
+                subfolder="scheduler",
+                torch_dtype=self.weights_dtype,
+                beta_start=0.00085,
+                beta_end=0.0120,
+                beta_schedule="scaled_linear",
+                cache_dir=self.cfg.cache_dir,
+            )
+        else:
+            self.scheduler = DDIMScheduler.from_pretrained(
+                self.cfg.pretrained_model_name_or_path,
+                subfolder="scheduler",
+                torch_dtype=self.weights_dtype,
+                cache_dir=self.cfg.cache_dir,
+                local_files_only=self.cfg.local_files_only,
+            )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+        if self.cfg.time_prior is not None:
+            m1, m2, s1, s2 = self.cfg.time_prior
+            weights = torch.cat(
+                (
+                    torch.exp(
+                        -((torch.arange(self.num_train_timesteps, m1, -1) - m1) ** 2)
+                        / (2 * s1**2)
+                    ),
+                    torch.ones(m1 - m2 + 1),
+                    torch.exp(
+                        -((torch.arange(m2 - 1, 0, -1) - m2) ** 2) / (2 * s2**2)
+                    ),
+                )
+            )
+            weights = weights / torch.sum(weights)
+            self.time_prior_acc_weights = torch.cumsum(weights, dim=0)
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        if self.cfg.use_sjc:
+            # score jacobian chaining need mu
+            self.us: Float[Tensor, "..."] = torch.sqrt((1 - self.alphas) / self.alphas)
+
+        self.grad_clip_val: Optional[float] = None
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        latent_height: int = 64,
+        latent_width: int = 64,
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        latents = F.interpolate(
+            latents, (latent_height, latent_width), mode="bilinear", align_corners=False
+        )
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def compute_grad_sds(
+        self,
+        latents: Float[Tensor, "B 4 64 64"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ):
+        batch_size = elevation.shape[0]
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        if self.cfg.weighting_strategy == "sds":
+            # w(t), sigma_t^2
+            w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (noise_pred - noise)
+
+        guidance_eval_utils = {
+            "use_perp_neg": prompt_utils.use_perp_neg,
+            "neg_guidance_weights": neg_guidance_weights,
+            "text_embeddings": text_embeddings,
+            "t_orig": t,
+            "latents_noisy": latents_noisy,
+            "noise_pred": noise_pred,
+        }
+
+        return grad, guidance_eval_utils
+
+    def compute_grad_sjc(
+        self,
+        latents: Float[Tensor, "B 4 64 64"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ):
+        batch_size = elevation.shape[0]
+
+        sigma = self.us[t]
+        sigma = sigma.view(-1, 1, 1, 1)
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                y = latents
+                zs = y + sigma * noise
+                scaled_zs = zs / torch.sqrt(1 + sigma**2)
+                # pred noise
+                latent_model_input = torch.cat([scaled_zs] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                y = latents
+
+                zs = y + sigma * noise
+                scaled_zs = zs / torch.sqrt(1 + sigma**2)
+
+                # pred noise
+                latent_model_input = torch.cat([scaled_zs] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )
+
+                # perform guidance (high scale from paper!)
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+        Ds = zs - sigma * noise_pred
+
+        if self.cfg.var_red:
+            grad = -(Ds - y) / sigma
+        else:
+            grad = -(Ds - zs) / sigma
+
+        guidance_eval_utils = {
+            "use_perp_neg": prompt_utils.use_perp_neg,
+            "neg_guidance_weights": neg_guidance_weights,
+            "text_embeddings": text_embeddings,
+            "t_orig": t,
+            "latents_noisy": scaled_zs,
+            "noise_pred": noise_pred,
+        }
+
+        return grad, guidance_eval_utils
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        guidance_eval=False,
+        current_step_ratio=None,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = F.interpolate(
+                rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        if self.cfg.time_prior is not None:
+            time_index = torch.where(
+                (self.time_prior_acc_weights - current_step_ratio) > 0
+            )[0][0]
+            if time_index == 0 or torch.abs(
+                self.time_prior_acc_weights[time_index] - current_step_ratio
+            ) < torch.abs(
+                self.time_prior_acc_weights[time_index - 1] - current_step_ratio
+            ):
+                t = self.num_train_timesteps - time_index
+            else:
+                t = self.num_train_timesteps - time_index + 1
+            t = torch.clip(t, self.min_step, self.max_step + 1)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+
+        else:
+            # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+            t = torch.randint(
+                self.min_step,
+                self.max_step + 1,
+                [batch_size],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+        if self.cfg.use_sjc:
+            grad, guidance_eval_utils = self.compute_grad_sjc(
+                latents, t, prompt_utils, elevation, azimuth, camera_distances
+            )
+        else:
+            grad, guidance_eval_utils = self.compute_grad_sds(
+                latents, t, prompt_utils, elevation, azimuth, camera_distances
+            )
+
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        if guidance_eval:
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_noise_pred(
+        self,
+        latents_noisy,
+        t,
+        text_embeddings,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        batch_size = latents_noisy.shape[0]
+
+        if use_perp_neg:
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 4).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 2).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return noise_pred
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(
+        self,
+        t_orig,
+        text_embeddings,
+        latents_noisy,
+        noise_pred,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = self.decode_latents(latents_noisy[:bs]).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1], eta=1
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = self.decode_latents(latents_1step).permute(0, 2, 3, 1)
+        imgs_1orig = self.decode_latents(pred_1orig).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            text_emb = (
+                text_embeddings[
+                    [b, b + len(idxs), b + 2 * len(idxs), b + 3 * len(idxs)], ...
+                ]
+                if use_perp_neg
+                else text_embeddings[[b, b + len(idxs)], ...]
+            )
+            neg_guid = neg_guidance_weights[b : b + 1] if use_perp_neg else None
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                noise_pred = self.get_noise_pred(
+                    latents, t, text_emb, use_perp_neg, neg_guid
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents, eta=1)[
+                    "prev_sample"
+                ]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = self.decode_latents(latents_final).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )

threestudio/models/guidance/stable_diffusion_unified_guidance.py

@@ -0,0 +1,729 @@
+import random
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import (
+    AutoencoderKL,
+    ControlNetModel,
+    DDPMScheduler,
+    DPMSolverSinglestepScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.networks import ToDTypeWrapper
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseModule
+from threestudio.utils.misc import C, cleanup, enable_gradient, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-unified-guidance")
+class StableDiffusionUnifiedGuidance(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+
+        # guidance type, in ["sds", "vsd"]
+        guidance_type: str = "sds"
+
+        pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        guidance_scale: float = 100.0
+        weighting_strategy: str = "dreamfusion"
+        view_dependent_prompting: bool = True
+
+        min_step_percent: Any = 0.02
+        max_step_percent: Any = 0.98
+        grad_clip: Optional[Any] = None
+
+        return_rgb_1step_orig: bool = False
+        return_rgb_multistep_orig: bool = False
+        n_rgb_multistep_orig_steps: int = 4
+
+        # TODO
+        # controlnet
+        controlnet_model_name_or_path: Optional[str] = None
+        preprocessor: Optional[str] = None
+        control_scale: float = 1.0
+
+        # TODO
+        # lora
+        lora_model_name_or_path: Optional[str] = None
+
+        # efficiency-related configurations
+        half_precision_weights: bool = True
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        token_merging: bool = False
+        token_merging_params: Optional[dict] = field(default_factory=dict)
+
+        # VSD configurations, only used when guidance_type is "vsd"
+        vsd_phi_model_name_or_path: Optional[str] = None
+        vsd_guidance_scale_phi: float = 1.0
+        vsd_use_lora: bool = True
+        vsd_lora_cfg_training: bool = False
+        vsd_lora_n_timestamp_samples: int = 1
+        vsd_use_camera_condition: bool = True
+        # camera condition type, in ["extrinsics", "mvp", "spherical"]
+        vsd_camera_condition_type: Optional[str] = "extrinsics"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.min_step: Optional[int] = None
+        self.max_step: Optional[int] = None
+        self.grad_clip_val: Optional[float] = None
+
+        @dataclass
+        class NonTrainableModules:
+            pipe: StableDiffusionPipeline
+            pipe_phi: Optional[StableDiffusionPipeline] = None
+            controlnet: Optional[ControlNetModel] = None
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        threestudio.info(f"Loading Stable Diffusion ...")
+
+        pipe_kwargs = {
+            "tokenizer": None,
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only,
+        }
+        pipe = StableDiffusionPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+        self.prepare_pipe(pipe)
+        self.configure_pipe_token_merging(pipe)
+
+        # phi network for VSD
+        # introduce two trainable modules:
+        # - self.camera_embedding
+        # - self.lora_layers
+        pipe_phi = None
+
+        # if the phi network shares the same unet with the pretrain network
+        # we need to pass additional cross attention kwargs to the unet
+        self.vsd_share_model = (
+            self.cfg.guidance_type == "vsd"
+            and self.cfg.vsd_phi_model_name_or_path is None
+        )
+        if self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_phi_model_name_or_path is None:
+                pipe_phi = pipe
+            else:
+                pipe_phi = StableDiffusionPipeline.from_pretrained(
+                    self.cfg.vsd_phi_model_name_or_path,
+                    **pipe_kwargs,
+                ).to(self.device)
+                self.prepare_pipe(pipe_phi)
+                self.configure_pipe_token_merging(pipe_phi)
+
+            # set up camera embedding
+            if self.cfg.vsd_use_camera_condition:
+                if self.cfg.vsd_camera_condition_type in ["extrinsics", "mvp"]:
+                    self.camera_embedding_dim = 16
+                elif self.cfg.vsd_camera_condition_type == "spherical":
+                    self.camera_embedding_dim = 4
+                else:
+                    raise ValueError("Invalid camera condition type!")
+
+                # FIXME: hard-coded output dim
+                self.camera_embedding = ToDTypeWrapper(
+                    TimestepEmbedding(self.camera_embedding_dim, 1280),
+                    self.weights_dtype,
+                ).to(self.device)
+                pipe_phi.unet.class_embedding = self.camera_embedding
+
+            if self.cfg.vsd_use_lora:
+                # set up LoRA layers
+                lora_attn_procs = {}
+                for name in pipe_phi.unet.attn_processors.keys():
+                    cross_attention_dim = (
+                        None
+                        if name.endswith("attn1.processor")
+                        else pipe_phi.unet.config.cross_attention_dim
+                    )
+                    if name.startswith("mid_block"):
+                        hidden_size = pipe_phi.unet.config.block_out_channels[-1]
+                    elif name.startswith("up_blocks"):
+                        block_id = int(name[len("up_blocks.")])
+                        hidden_size = list(
+                            reversed(pipe_phi.unet.config.block_out_channels)
+                        )[block_id]
+                    elif name.startswith("down_blocks"):
+                        block_id = int(name[len("down_blocks.")])
+                        hidden_size = pipe_phi.unet.config.block_out_channels[block_id]
+
+                    lora_attn_procs[name] = LoRAAttnProcessor(
+                        hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
+                    )
+
+                pipe_phi.unet.set_attn_processor(lora_attn_procs)
+
+                self.lora_layers = AttnProcsLayers(pipe_phi.unet.attn_processors).to(
+                    self.device
+                )
+                self.lora_layers._load_state_dict_pre_hooks.clear()
+                self.lora_layers._state_dict_hooks.clear()
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+        # controlnet
+        controlnet = None
+        if self.cfg.controlnet_model_name_or_path is not None:
+            threestudio.info(f"Loading ControlNet ...")
+
+            controlnet = ControlNetModel.from_pretrained(
+                self.cfg.controlnet_model_name_or_path,
+                torch_dtype=self.weights_dtype,
+            ).to(self.device)
+            controlnet.eval()
+            enable_gradient(controlnet, enabled=False)
+
+            threestudio.info(f"Loaded ControlNet!")
+
+        self.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        # q(z_t|x) = N(alpha_t x, sigma_t^2 I)
+        # in DDPM, alpha_t = sqrt(alphas_cumprod_t), sigma_t^2 = 1 - alphas_cumprod_t
+        self.alphas_cumprod: Float[Tensor, "T"] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        self.alphas: Float[Tensor, "T"] = self.alphas_cumprod**0.5
+        self.sigmas: Float[Tensor, "T"] = (1 - self.alphas_cumprod) ** 0.5
+        # log SNR
+        self.lambdas: Float[Tensor, "T"] = self.sigmas / self.alphas
+
+        self._non_trainable_modules = NonTrainableModules(
+            pipe=pipe,
+            pipe_phi=pipe_phi,
+            controlnet=controlnet,
+        )
+
+    @property
+    def pipe(self) -> StableDiffusionPipeline:
+        return self._non_trainable_modules.pipe
+
+    @property
+    def pipe_phi(self) -> StableDiffusionPipeline:
+        if self._non_trainable_modules.pipe_phi is None:
+            raise RuntimeError("phi model is not available.")
+        return self._non_trainable_modules.pipe_phi
+
+    @property
+    def controlnet(self) -> ControlNetModel:
+        if self._non_trainable_modules.controlnet is None:
+            raise RuntimeError("ControlNet model is not available.")
+        return self._non_trainable_modules.controlnet
+
+    def prepare_pipe(self, pipe: StableDiffusionPipeline):
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            pipe.unet.to(memory_format=torch.channels_last)
+
+        # FIXME: pipe.__call__ requires text_encoder.dtype
+        # pipe.text_encoder.to("meta")
+        cleanup()
+
+        pipe.vae.eval()
+        pipe.unet.eval()
+
+        enable_gradient(pipe.vae, enabled=False)
+        enable_gradient(pipe.unet, enabled=False)
+
+        # disable progress bar
+        pipe.set_progress_bar_config(disable=True)
+
+    def configure_pipe_token_merging(self, pipe: StableDiffusionPipeline):
+        if self.cfg.token_merging:
+            import tomesd
+
+            tomesd.apply_patch(pipe.unet, **self.cfg.token_merging_params)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        unet: UNet2DConditionModel,
+        latents: Float[Tensor, "..."],
+        t: Int[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        class_labels: Optional[Float[Tensor, "..."]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Float[Tensor, "..."]] = None,
+        mid_block_additional_residual: Optional[Float[Tensor, "..."]] = None,
+        velocity_to_epsilon: bool = False,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        pred = unet(
+            latents.to(unet.dtype),
+            t.to(unet.dtype),
+            encoder_hidden_states=encoder_hidden_states.to(unet.dtype),
+            class_labels=class_labels,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample
+        if velocity_to_epsilon:
+            pred = latents * self.sigmas[t].view(-1, 1, 1, 1) + pred * self.alphas[
+                t
+            ].view(-1, 1, 1, 1)
+        return pred.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_encode(
+        self, vae: AutoencoderKL, imgs: Float[Tensor, "B 3 H W"], mode=False
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        # expect input in [-1, 1]
+        input_dtype = imgs.dtype
+        posterior = vae.encode(imgs.to(vae.dtype)).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample()
+        latents = latents * vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_decode(
+        self, vae: AutoencoderKL, latents: Float[Tensor, "B 4 Hl Wl"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        # output in [0, 1]
+        input_dtype = latents.dtype
+        latents = 1 / vae.config.scaling_factor * latents
+        image = vae.decode(latents.to(vae.dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @contextmanager
+    def disable_unet_class_embedding(self, unet: UNet2DConditionModel):
+        class_embedding = unet.class_embedding
+        try:
+            unet.class_embedding = None
+            yield unet
+        finally:
+            unet.class_embedding = class_embedding
+
+    @contextmanager
+    def set_scheduler(
+        self, pipe: StableDiffusionPipeline, scheduler_class: Any, **kwargs
+    ):
+        scheduler_orig = pipe.scheduler
+        pipe.scheduler = scheduler_class.from_config(scheduler_orig.config, **kwargs)
+        yield pipe
+        pipe.scheduler = scheduler_orig
+
+    def get_eps_pretrain(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                    noise_pred = self.forward_unet(
+                        unet,
+                        torch.cat([latents_noisy] * 4, dim=0),
+                        torch.cat([t] * 4, dim=0),
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                        == "v_prediction",
+                    )  # (4B, 3, Hl, Wl)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                    noise_pred = self.forward_unet(
+                        unet,
+                        torch.cat([latents_noisy] * 2, dim=0),
+                        torch.cat([t] * 2, dim=0),
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                        == "v_prediction",
+                    )
+
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return noise_pred
+
+    def get_eps_phi(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        # not using view-dependent prompting in LoRA
+        text_embeddings, _ = prompt_utils.get_text_embeddings(
+            elevation, azimuth, camera_distances, view_dependent_prompting=False
+        ).chunk(2)
+        with torch.no_grad():
+            noise_pred = self.forward_unet(
+                self.pipe_phi.unet,
+                torch.cat([latents_noisy] * 2, dim=0),
+                torch.cat([t] * 2, dim=0),
+                encoder_hidden_states=torch.cat([text_embeddings] * 2, dim=0),
+                class_labels=torch.cat(
+                    [
+                        camera_condition.view(batch_size, -1),
+                        torch.zeros_like(camera_condition.view(batch_size, -1)),
+                    ],
+                    dim=0,
+                )
+                if self.cfg.vsd_use_camera_condition
+                else None,
+                cross_attention_kwargs={"scale": 1.0},
+                velocity_to_epsilon=self.pipe_phi.scheduler.config.prediction_type
+                == "v_prediction",
+            )
+
+        noise_pred_camera, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.vsd_guidance_scale_phi * (
+            noise_pred_camera - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def train_phi(
+        self,
+        latents: Float[Tensor, "B 4 Hl Wl"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ):
+        B = latents.shape[0]
+        latents = latents.detach().repeat(
+            self.cfg.vsd_lora_n_timestamp_samples, 1, 1, 1
+        )
+
+        num_train_timesteps = self.pipe_phi.scheduler.config.num_train_timesteps
+        t = torch.randint(
+            int(num_train_timesteps * 0.0),
+            int(num_train_timesteps * 1.0),
+            [B * self.cfg.vsd_lora_n_timestamp_samples],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        noise = torch.randn_like(latents)
+        latents_noisy = self.pipe_phi.scheduler.add_noise(latents, noise, t)
+        if self.pipe_phi.scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif self.pipe_phi.scheduler.prediction_type == "v_prediction":
+            target = self.pipe_phi.scheduler.get_velocity(latents, noise, t)
+        else:
+            raise ValueError(
+                f"Unknown prediction type {self.pipe_phi.scheduler.prediction_type}"
+            )
+
+        # not using view-dependent prompting in LoRA
+        text_embeddings, _ = prompt_utils.get_text_embeddings(
+            elevation, azimuth, camera_distances, view_dependent_prompting=False
+        ).chunk(2)
+
+        if (
+            self.cfg.vsd_use_camera_condition
+            and self.cfg.vsd_lora_cfg_training
+            and random.random() < 0.1
+        ):
+            camera_condition = torch.zeros_like(camera_condition)
+
+        noise_pred = self.forward_unet(
+            self.pipe_phi.unet,
+            latents_noisy,
+            t,
+            encoder_hidden_states=text_embeddings.repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1, 1
+            ),
+            class_labels=camera_condition.view(B, -1).repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1
+            )
+            if self.cfg.vsd_use_camera_condition
+            else None,
+            cross_attention_kwargs={"scale": 1.0},
+        )
+        return F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+    def forward(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        mvp_mtx: Float[Tensor, "B 4 4"],
+        c2w: Float[Tensor, "B 4 4"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 Hl Wl"]
+        if rgb_as_latents:
+            # treat input rgb as latents
+            # input rgb should be in range [-1, 1]
+            latents = F.interpolate(
+                rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+            )
+        else:
+            # treat input rgb as rgb
+            # input rgb should be in range [0, 1]
+            rgb_BCHW = F.interpolate(
+                rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.vae_encode(self.pipe.vae, rgb_BCHW * 2.0 - 1.0)
+
+        # sample timestep
+        # use the same timestep for each batch
+        assert self.min_step is not None and self.max_step is not None
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [1],
+            dtype=torch.long,
+            device=self.device,
+        ).repeat(batch_size)
+
+        # sample noise
+        noise = torch.randn_like(latents)
+        latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+        eps_pretrain = self.get_eps_pretrain(
+            latents_noisy, t, prompt_utils, elevation, azimuth, camera_distances
+        )
+
+        latents_1step_orig = (
+            1
+            / self.alphas[t].view(-1, 1, 1, 1)
+            * (latents_noisy - self.sigmas[t].view(-1, 1, 1, 1) * eps_pretrain)
+        ).detach()
+
+        if self.cfg.guidance_type == "sds":
+            eps_phi = noise
+        elif self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_camera_condition_type == "extrinsics":
+                camera_condition = c2w
+            elif self.cfg.vsd_camera_condition_type == "mvp":
+                camera_condition = mvp_mtx
+            elif self.cfg.vsd_camera_condition_type == "spherical":
+                camera_condition = torch.stack(
+                    [
+                        torch.deg2rad(elevation),
+                        torch.sin(torch.deg2rad(azimuth)),
+                        torch.cos(torch.deg2rad(azimuth)),
+                        camera_distances,
+                    ],
+                    dim=-1,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown camera_condition_type {self.cfg.vsd_camera_condition_type}"
+                )
+            eps_phi = self.get_eps_phi(
+                latents_noisy,
+                t,
+                prompt_utils,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+            loss_train_phi = self.train_phi(
+                latents,
+                prompt_utils,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+        if self.cfg.weighting_strategy == "dreamfusion":
+            w = (1.0 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1.0
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (eps_pretrain - eps_phi)
+
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # reparameterization trick:
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        target = (latents - grad).detach()
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "timesteps": t,
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+            "latents": latents,
+            "latents_1step_orig": latents_1step_orig,
+            "rgb": rgb_BCHW.permute(0, 2, 3, 1),
+            "weights": w,
+            "lambdas": self.lambdas[t],
+        }
+
+        if self.cfg.return_rgb_1step_orig:
+            with torch.no_grad():
+                rgb_1step_orig = self.vae_decode(
+                    self.pipe.vae, latents_1step_orig
+                ).permute(0, 2, 3, 1)
+            guidance_out.update({"rgb_1step_orig": rgb_1step_orig})
+
+        if self.cfg.return_rgb_multistep_orig:
+            with self.set_scheduler(
+                self.pipe,
+                DPMSolverSinglestepScheduler,
+                solver_order=1,
+                num_train_timesteps=int(t[0]),
+            ) as pipe:
+                text_embeddings = prompt_utils.get_text_embeddings(
+                    elevation,
+                    azimuth,
+                    camera_distances,
+                    self.cfg.view_dependent_prompting,
+                )
+                text_embeddings_cond, text_embeddings_uncond = text_embeddings.chunk(2)
+                with torch.cuda.amp.autocast(enabled=False):
+                    latents_multistep_orig = pipe(
+                        num_inference_steps=self.cfg.n_rgb_multistep_orig_steps,
+                        guidance_scale=self.cfg.guidance_scale,
+                        eta=1.0,
+                        latents=latents_noisy.to(pipe.unet.dtype),
+                        prompt_embeds=text_embeddings_cond.to(pipe.unet.dtype),
+                        negative_prompt_embeds=text_embeddings_uncond.to(
+                            pipe.unet.dtype
+                        ),
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        output_type="latent",
+                    ).images.to(latents.dtype)
+            with torch.no_grad():
+                rgb_multistep_orig = self.vae_decode(
+                    self.pipe.vae, latents_multistep_orig
+                )
+            guidance_out.update(
+                {
+                    "latents_multistep_orig": latents_multistep_orig,
+                    "rgb_multistep_orig": rgb_multistep_orig.permute(0, 2, 3, 1),
+                }
+            )
+
+        if self.cfg.guidance_type == "vsd":
+            guidance_out.update(
+                {
+                    "loss_train_phi": loss_train_phi,
+                }
+            )
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.min_step = int(
+            self.num_train_timesteps * C(self.cfg.min_step_percent, epoch, global_step)
+        )
+        self.max_step = int(
+            self.num_train_timesteps * C(self.cfg.max_step_percent, epoch, global_step)
+        )

threestudio/models/guidance/stable_zero123_guidance.py

@@ -0,0 +1,340 @@
+import importlib
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == "__is_first_stage__":
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+# load model
+def load_model_from_config(config, ckpt, device, vram_O=True, verbose=False):
+    pl_sd = torch.load(ckpt, map_location="cpu")
+
+    if "global_step" in pl_sd and verbose:
+        print(f'[INFO] Global Step: {pl_sd["global_step"]}')
+
+    sd = pl_sd["state_dict"]
+
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+
+    if len(m) > 0 and verbose:
+        print("[INFO] missing keys: \n", m)
+    if len(u) > 0 and verbose:
+        print("[INFO] unexpected keys: \n", u)
+
+    # manually load ema and delete it to save GPU memory
+    if model.use_ema:
+        if verbose:
+            print("[INFO] loading EMA...")
+        model.model_ema.copy_to(model.model)
+        del model.model_ema
+
+    if vram_O:
+        # we don't need decoder
+        del model.first_stage_model.decoder
+
+    torch.cuda.empty_cache()
+
+    model.eval().to(device)
+
+    return model
+
+
+@threestudio.register("stable-zero123-guidance")
+class StableZero123Guidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        pretrained_model_name_or_path: str = "load/zero123/stable-zero123.ckpt"
+        pretrained_config: str = "load/zero123/sd-objaverse-finetune-c_concat-256.yaml"
+        vram_O: bool = True
+
+        cond_image_path: str = "load/images/hamburger_rgba.png"
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        guidance_scale: float = 5.0
+
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        half_precision_weights: bool = False
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Stable Zero123 ...")
+
+        self.config = OmegaConf.load(self.cfg.pretrained_config)
+        # TODO: seems it cannot load into fp16...
+        self.weights_dtype = torch.float32
+        self.model = load_model_from_config(
+            self.config,
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            vram_O=self.cfg.vram_O,
+        )
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        # timesteps: use diffuser for convenience... hope it's alright.
+        self.num_train_timesteps = self.config.model.params.timesteps
+
+        self.scheduler = DDIMScheduler(
+            self.num_train_timesteps,
+            self.config.model.params.linear_start,
+            self.config.model.params.linear_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        self.prepare_embeddings(self.cfg.cond_image_path)
+
+        threestudio.info(f"Loaded Stable Zero123!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def prepare_embeddings(self, image_path: str) -> None:
+        # load cond image for zero123
+        assert os.path.exists(image_path)
+        rgba = cv2.cvtColor(
+            cv2.imread(image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(rgba, (256, 256), interpolation=cv2.INTER_AREA).astype(
+                np.float32
+            )
+            / 255.0
+        )
+        rgb = rgba[..., :3] * rgba[..., 3:] + (1 - rgba[..., 3:])
+        self.rgb_256: Float[Tensor, "1 3 H W"] = (
+            torch.from_numpy(rgb)
+            .unsqueeze(0)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+            .to(self.device)
+        )
+        self.c_crossattn, self.c_concat = self.get_img_embeds(self.rgb_256)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_img_embeds(
+        self,
+        img: Float[Tensor, "B 3 256 256"],
+    ) -> Tuple[Float[Tensor, "B 1 768"], Float[Tensor, "B 4 32 32"]]:
+        img = img * 2.0 - 1.0
+        c_crossattn = self.model.get_learned_conditioning(img.to(self.weights_dtype))
+        c_concat = self.model.encode_first_stage(img.to(self.weights_dtype)).mode()
+        return c_crossattn, c_concat
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 256 256"]
+    ) -> Float[Tensor, "B 4 32 32"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        latents = self.model.get_first_stage_encoding(
+            self.model.encode_first_stage(imgs.to(self.weights_dtype))
+        )
+        return latents.to(input_dtype)  # [B, 4, 32, 32] Latent space image
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        image = self.model.decode_first_stage(latents)
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_cond(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        c_crossattn=None,
+        c_concat=None,
+        **kwargs,
+    ) -> dict:
+        T = torch.stack(
+            [
+                torch.deg2rad(
+                    (90 - elevation) - (90 - self.cfg.cond_elevation_deg)
+                ),  # Zero123 polar is 90-elevation
+                torch.sin(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.deg2rad(
+                    90 - torch.full_like(elevation, self.cfg.cond_elevation_deg)
+                ),
+            ],
+            dim=-1,
+        )[:, None, :].to(self.device)
+        cond = {}
+        clip_emb = self.model.cc_projection(
+            torch.cat(
+                [
+                    (self.c_crossattn if c_crossattn is None else c_crossattn).repeat(
+                        len(T), 1, 1
+                    ),
+                    T,
+                ],
+                dim=-1,
+            )
+        )
+        cond["c_crossattn"] = [
+            torch.cat([torch.zeros_like(clip_emb).to(self.device), clip_emb], dim=0)
+        ]
+        cond["c_concat"] = [
+            torch.cat(
+                [
+                    torch.zeros_like(self.c_concat)
+                    .repeat(len(T), 1, 1, 1)
+                    .to(self.device),
+                    (self.c_concat if c_concat is None else c_concat).repeat(
+                        len(T), 1, 1, 1
+                    ),
+                ],
+                dim=0,
+            )
+        ]
+        return cond
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = (
+                F.interpolate(rgb_BCHW, (32, 32), mode="bilinear", align_corners=False)
+                * 2
+                - 1
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        cond = self.get_cond(elevation, azimuth, camera_distances)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_cond - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).reshape(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )

threestudio/models/guidance/zero123_guidance.py

@@ -0,0 +1,528 @@
+import importlib
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == "__is_first_stage__":
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+# load model
+def load_model_from_config(config, ckpt, device, vram_O=True, verbose=False):
+    pl_sd = torch.load(ckpt, map_location="cpu")
+
+    if "global_step" in pl_sd and verbose:
+        print(f'[INFO] Global Step: {pl_sd["global_step"]}')
+
+    sd = pl_sd["state_dict"]
+
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+
+    if len(m) > 0 and verbose:
+        print("[INFO] missing keys: \n", m)
+    if len(u) > 0 and verbose:
+        print("[INFO] unexpected keys: \n", u)
+
+    # manually load ema and delete it to save GPU memory
+    if model.use_ema:
+        if verbose:
+            print("[INFO] loading EMA...")
+        model.model_ema.copy_to(model.model)
+        del model.model_ema
+
+    if vram_O:
+        # we don't need decoder
+        del model.first_stage_model.decoder
+
+    torch.cuda.empty_cache()
+
+    model.eval().to(device)
+
+    return model
+
+
+@threestudio.register("zero123-guidance")
+class Zero123Guidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        pretrained_model_name_or_path: str = "load/zero123/105000.ckpt"
+        pretrained_config: str = "load/zero123/sd-objaverse-finetune-c_concat-256.yaml"
+        vram_O: bool = True
+
+        cond_image_path: str = "load/images/hamburger_rgba.png"
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        guidance_scale: float = 5.0
+
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        half_precision_weights: bool = False
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Zero123 ...")
+
+        self.config = OmegaConf.load(self.cfg.pretrained_config)
+        # TODO: seems it cannot load into fp16...
+        self.weights_dtype = torch.float32
+        self.model = load_model_from_config(
+            self.config,
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            vram_O=self.cfg.vram_O,
+        )
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        # timesteps: use diffuser for convenience... hope it's alright.
+        self.num_train_timesteps = self.config.model.params.timesteps
+
+        self.scheduler = DDIMScheduler(
+            self.num_train_timesteps,
+            self.config.model.params.linear_start,
+            self.config.model.params.linear_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        self.prepare_embeddings(self.cfg.cond_image_path)
+
+        threestudio.info(f"Loaded Zero123!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def prepare_embeddings(self, image_path: str) -> None:
+        # load cond image for zero123
+        assert os.path.exists(image_path)
+        rgba = cv2.cvtColor(
+            cv2.imread(image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(rgba, (256, 256), interpolation=cv2.INTER_AREA).astype(
+                np.float32
+            )
+            / 255.0
+        )
+        rgb = rgba[..., :3] * rgba[..., 3:] + (1 - rgba[..., 3:])
+        self.rgb_256: Float[Tensor, "1 3 H W"] = (
+            torch.from_numpy(rgb)
+            .unsqueeze(0)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+            .to(self.device)
+        )
+        self.c_crossattn, self.c_concat = self.get_img_embeds(self.rgb_256)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_img_embeds(
+        self,
+        img: Float[Tensor, "B 3 256 256"],
+    ) -> Tuple[Float[Tensor, "B 1 768"], Float[Tensor, "B 4 32 32"]]:
+        img = img * 2.0 - 1.0
+        c_crossattn = self.model.get_learned_conditioning(img.to(self.weights_dtype))
+        c_concat = self.model.encode_first_stage(img.to(self.weights_dtype)).mode()
+        return c_crossattn, c_concat
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 256 256"]
+    ) -> Float[Tensor, "B 4 32 32"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        latents = self.model.get_first_stage_encoding(
+            self.model.encode_first_stage(imgs.to(self.weights_dtype))
+        )
+        return latents.to(input_dtype)  # [B, 4, 32, 32] Latent space image
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        image = self.model.decode_first_stage(latents)
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_cond(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        c_crossattn=None,
+        c_concat=None,
+        **kwargs,
+    ) -> dict:
+        T = torch.stack(
+            [
+                torch.deg2rad(
+                    (90 - elevation) - (90 - self.cfg.cond_elevation_deg)
+                ),  # Zero123 polar is 90-elevation
+                torch.sin(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                camera_distances - self.cfg.cond_camera_distance,
+            ],
+            dim=-1,
+        )[:, None, :].to(self.device)
+        cond = {}
+        clip_emb = self.model.cc_projection(
+            torch.cat(
+                [
+                    (self.c_crossattn if c_crossattn is None else c_crossattn).repeat(
+                        len(T), 1, 1
+                    ),
+                    T,
+                ],
+                dim=-1,
+            )
+        )
+        cond["c_crossattn"] = [
+            torch.cat([torch.zeros_like(clip_emb).to(self.device), clip_emb], dim=0)
+        ]
+        cond["c_concat"] = [
+            torch.cat(
+                [
+                    torch.zeros_like(self.c_concat)
+                    .repeat(len(T), 1, 1, 1)
+                    .to(self.device),
+                    (self.c_concat if c_concat is None else c_concat).repeat(
+                        len(T), 1, 1, 1
+                    ),
+                ],
+                dim=0,
+            )
+        ]
+        return cond
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        guidance_eval=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = (
+                F.interpolate(rgb_BCHW, (32, 32), mode="bilinear", align_corners=False)
+                * 2
+                - 1
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        cond = self.get_cond(elevation, azimuth, camera_distances)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_cond - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).reshape(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds, # loss_sds
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        if guidance_eval:
+            guidance_eval_utils = {
+                "cond": cond,
+                "t_orig": t,
+                "latents_noisy": latents_noisy,
+                "noise_pred": noise_pred,
+            }
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(self, cond, t_orig, latents_noisy, noise_pred):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = self.decode_latents(latents_noisy[:bs]).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1], eta=1
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = self.decode_latents(latents_1step).permute(0, 2, 3, 1)
+        imgs_1orig = self.decode_latents(pred_1orig).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            c = {
+                "c_crossattn": [cond["c_crossattn"][0][[b, b + len(idxs)], ...]],
+                "c_concat": [cond["c_concat"][0][[b, b + len(idxs)], ...]],
+            }
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                x_in = torch.cat([latents] * 2)
+                t_in = torch.cat([t.reshape(1)] * 2).to(self.device)
+                noise_pred = self.model.apply_model(x_in, t_in, c)
+                # perform guidance
+                noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_cond - noise_pred_uncond
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents, eta=1)[
+                    "prev_sample"
+                ]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = self.decode_latents(latents_final).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+    # verification - requires `vram_O = False` in load_model_from_config
+    @torch.no_grad()
+    def generate(
+        self,
+        image,  # image tensor [1, 3, H, W] in [0, 1]
+        elevation=0,
+        azimuth=0,
+        camera_distances=0,  # new view params
+        c_crossattn=None,
+        c_concat=None,
+        scale=3,
+        ddim_steps=50,
+        post_process=True,
+        ddim_eta=1,
+    ):
+        if c_crossattn is None:
+            c_crossattn, c_concat = self.get_img_embeds(image)
+
+        cond = self.get_cond(
+            elevation, azimuth, camera_distances, c_crossattn, c_concat
+        )
+
+        imgs = self.gen_from_cond(cond, scale, ddim_steps, post_process, ddim_eta)
+
+        return imgs
+
+    # verification - requires `vram_O = False` in load_model_from_config
+    @torch.no_grad()
+    def gen_from_cond(
+        self,
+        cond,
+        scale=3,
+        ddim_steps=50,
+        post_process=True,
+        ddim_eta=1,
+    ):
+        # produce latents loop
+        B = cond["c_crossattn"][0].shape[0] // 2
+        latents = torch.randn((B, 4, 32, 32), device=self.device)
+        self.scheduler.set_timesteps(ddim_steps)
+
+        for t in self.scheduler.timesteps:
+            x_in = torch.cat([latents] * 2)
+            t_in = torch.cat([t.reshape(1).repeat(B)] * 2).to(self.device)
+
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + scale * (
+                noise_pred_cond - noise_pred_uncond
+            )
+
+            latents = self.scheduler.step(noise_pred, t, latents, eta=ddim_eta)[
+                "prev_sample"
+            ]
+
+        imgs = self.decode_latents(latents)
+        imgs = imgs.cpu().numpy().transpose(0, 2, 3, 1) if post_process else imgs
+
+        return imgs
+
+
+if __name__ == '__main__':
+    from threestudio.utils.config import load_config
+    import pytorch_lightning as pl
+    import numpy as np
+    import os
+    import cv2
+    cfg = load_config("configs/experimental/zero123.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    elevations = [0, 20, -20]
+    azimuths = [45, 90, 135, -45, -90]
+    radius = torch.tensor([3.8]).to(guidance.device)
+    outdir = ".threestudio_cache/saiyan"
+    os.makedirs(outdir, exist_ok=True)
+    # rgb_image = (rgb_image[0].detach().cpu().clip(0, 1).numpy()*255).astype(np.uint8)[:, :, ::-1].copy()
+    # os.makedirs('.threestudio_cache', exist_ok=True)
+    # cv2.imwrite('.threestudio_cache/diffusion_image.jpg', rgb_image)
+
+
+    rgb_image = cv2.imread(cfg.system.guidance.cond_image_path)[:, :, ::-1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (256, 256))
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device).permute(0,3,1,2)
+
+    for elevation in elevations:
+        for azimuth in azimuths:
+            output1 = guidance.generate(
+                rgb_image,
+                torch.tensor([elevation]).to(guidance.device),
+                torch.tensor([azimuth]).to(guidance.device),
+                radius,
+                c_crossattn=guidance.c_crossattn,
+                c_concat=guidance.c_concat
+            )
+            from torchvision.utils import save_image
+            save_image(torch.tensor(output1).float().permute(0,3,1,2), f"{outdir}/result_e_{elevation}_a_{azimuth}.png", normalize=True, value_range=(0,1))
+            

threestudio/models/guidance/zero123_unified_guidance.py

@@ -0,0 +1,721 @@
+import os
+import random
+import sys
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DPMSolverSinglestepScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding
+from PIL import Image
+from tqdm import tqdm
+
+import threestudio
+from extern.zero123 import Zero123Pipeline
+from threestudio.models.networks import ToDTypeWrapper
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseModule
+from threestudio.utils.misc import C, cleanup, enable_gradient, parse_version
+from threestudio.utils.typing import *
+
+
+@threestudio.register("zero123-unified-guidance")
+class Zero123UnifiedGuidance(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+
+        # guidance type, in ["sds", "vsd"]
+        guidance_type: str = "sds"
+
+        pretrained_model_name_or_path: str = "bennyguo/zero123-diffusers"
+        guidance_scale: float = 5.0
+        weighting_strategy: str = "dreamfusion"
+
+        min_step_percent: Any = 0.02
+        max_step_percent: Any = 0.98
+        grad_clip: Optional[Any] = None
+
+        return_rgb_1step_orig: bool = False
+        return_rgb_multistep_orig: bool = False
+        n_rgb_multistep_orig_steps: int = 4
+
+        cond_image_path: str = ""
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        # efficiency-related configurations
+        half_precision_weights: bool = True
+
+        # VSD configurations, only used when guidance_type is "vsd"
+        vsd_phi_model_name_or_path: Optional[str] = None
+        vsd_guidance_scale_phi: float = 1.0
+        vsd_use_lora: bool = True
+        vsd_lora_cfg_training: bool = False
+        vsd_lora_n_timestamp_samples: int = 1
+        vsd_use_camera_condition: bool = True
+        # camera condition type, in ["extrinsics", "mvp", "spherical"]
+        vsd_camera_condition_type: Optional[str] = "extrinsics"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.min_step: Optional[int] = None
+        self.max_step: Optional[int] = None
+        self.grad_clip_val: Optional[float] = None
+
+        @dataclass
+        class NonTrainableModules:
+            pipe: Zero123Pipeline
+            pipe_phi: Optional[Zero123Pipeline] = None
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        threestudio.info(f"Loading Zero123 ...")
+
+        # need to make sure the pipeline file is in path
+        sys.path.append("extern/")
+
+        pipe_kwargs = {
+            "safety_checker": None,
+            "requires_safety_checker": False,
+            "variant": "fp16" if self.cfg.half_precision_weights else None,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only,
+        }
+        pipe = Zero123Pipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+        self.prepare_pipe(pipe)
+
+        # phi network for VSD
+        # introduce two trainable modules:
+        # - self.camera_embedding
+        # - self.lora_layers
+        pipe_phi = None
+
+        # if the phi network shares the same unet with the pretrain network
+        # we need to pass additional cross attention kwargs to the unet
+        self.vsd_share_model = (
+            self.cfg.guidance_type == "vsd"
+            and self.cfg.vsd_phi_model_name_or_path is None
+        )
+        if self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_phi_model_name_or_path is None:
+                pipe_phi = pipe
+            else:
+                pipe_phi = Zero123Pipeline.from_pretrained(
+                    self.cfg.vsd_phi_model_name_or_path,
+                    **pipe_kwargs,
+                ).to(self.device)
+                self.prepare_pipe(pipe_phi)
+
+            # set up camera embedding
+            if self.cfg.vsd_use_camera_condition:
+                if self.cfg.vsd_camera_condition_type in ["extrinsics", "mvp"]:
+                    self.camera_embedding_dim = 16
+                elif self.cfg.vsd_camera_condition_type == "spherical":
+                    self.camera_embedding_dim = 4
+                else:
+                    raise ValueError("Invalid camera condition type!")
+
+                # FIXME: hard-coded output dim
+                self.camera_embedding = ToDTypeWrapper(
+                    TimestepEmbedding(self.camera_embedding_dim, 1280),
+                    self.weights_dtype,
+                ).to(self.device)
+                pipe_phi.unet.class_embedding = self.camera_embedding
+
+            if self.cfg.vsd_use_lora:
+                # set up LoRA layers
+                lora_attn_procs = {}
+                for name in pipe_phi.unet.attn_processors.keys():
+                    cross_attention_dim = (
+                        None
+                        if name.endswith("attn1.processor")
+                        else pipe_phi.unet.config.cross_attention_dim
+                    )
+                    if name.startswith("mid_block"):
+                        hidden_size = pipe_phi.unet.config.block_out_channels[-1]
+                    elif name.startswith("up_blocks"):
+                        block_id = int(name[len("up_blocks.")])
+                        hidden_size = list(
+                            reversed(pipe_phi.unet.config.block_out_channels)
+                        )[block_id]
+                    elif name.startswith("down_blocks"):
+                        block_id = int(name[len("down_blocks.")])
+                        hidden_size = pipe_phi.unet.config.block_out_channels[block_id]
+
+                    lora_attn_procs[name] = LoRAAttnProcessor(
+                        hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
+                    )
+
+                pipe_phi.unet.set_attn_processor(lora_attn_procs)
+
+                self.lora_layers = AttnProcsLayers(pipe_phi.unet.attn_processors).to(
+                    self.device
+                )
+                self.lora_layers._load_state_dict_pre_hooks.clear()
+                self.lora_layers._state_dict_hooks.clear()
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+        self.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        # q(z_t|x) = N(alpha_t x, sigma_t^2 I)
+        # in DDPM, alpha_t = sqrt(alphas_cumprod_t), sigma_t^2 = 1 - alphas_cumprod_t
+        self.alphas_cumprod: Float[Tensor, "T"] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        self.alphas: Float[Tensor, "T"] = self.alphas_cumprod**0.5
+        self.sigmas: Float[Tensor, "T"] = (1 - self.alphas_cumprod) ** 0.5
+        # log SNR
+        self.lambdas: Float[Tensor, "T"] = self.sigmas / self.alphas
+
+        self._non_trainable_modules = NonTrainableModules(
+            pipe=pipe,
+            pipe_phi=pipe_phi,
+        )
+
+        # self.clip_image_embeddings and self.image_latents
+        self.prepare_image_embeddings()
+
+    @property
+    def pipe(self) -> Zero123Pipeline:
+        return self._non_trainable_modules.pipe
+
+    @property
+    def pipe_phi(self) -> Zero123Pipeline:
+        if self._non_trainable_modules.pipe_phi is None:
+            raise RuntimeError("phi model is not available.")
+        return self._non_trainable_modules.pipe_phi
+
+    def prepare_pipe(self, pipe: Zero123Pipeline):
+        cleanup()
+
+        pipe.image_encoder.eval()
+        pipe.vae.eval()
+        pipe.unet.eval()
+        pipe.clip_camera_projection.eval()
+
+        enable_gradient(pipe.image_encoder, enabled=False)
+        enable_gradient(pipe.vae, enabled=False)
+        enable_gradient(pipe.unet, enabled=False)
+        enable_gradient(pipe.clip_camera_projection, enabled=False)
+
+        # disable progress bar
+        pipe.set_progress_bar_config(disable=True)
+
+    def prepare_image_embeddings(self) -> None:
+        if not os.path.exists(self.cfg.cond_image_path):
+            raise RuntimeError(
+                f"Condition image not found at {self.cfg.cond_image_path}"
+            )
+        image = Image.open(self.cfg.cond_image_path).convert("RGBA").resize((256, 256))
+        image = (
+            TF.to_tensor(image)
+            .unsqueeze(0)
+            .to(device=self.device, dtype=self.weights_dtype)
+        )
+        # rgba -> rgb, apply white background
+        image = image[:, :3] * image[:, 3:4] + (1 - image[:, 3:4])
+
+        with torch.no_grad():
+            self.clip_image_embeddings: Float[
+                Tensor, "1 1 D"
+            ] = self.extract_clip_image_embeddings(image)
+
+            # encoded latents should be multiplied with vae.config.scaling_factor
+            # but zero123 was not trained this way
+            self.image_latents: Float[Tensor, "1 4 Hl Wl"] = (
+                self.vae_encode(self.pipe.vae, image * 2.0 - 1.0, mode=True)
+                / self.pipe.vae.config.scaling_factor
+            )
+
+    def extract_clip_image_embeddings(
+        self, images: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 1 D"]:
+        # expect images in [0, 1]
+        images_pil = [TF.to_pil_image(image) for image in images]
+        images_processed = self.pipe.feature_extractor(
+            images=images_pil, return_tensors="pt"
+        ).pixel_values.to(device=self.device, dtype=self.weights_dtype)
+        clip_image_embeddings = self.pipe.image_encoder(images_processed).image_embeds
+        return clip_image_embeddings.to(images.dtype)
+
+    def get_image_camera_embeddings(
+        self,
+        elevation_deg: Float[Tensor, "B"],
+        azimuth_deg: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 1 D"]:
+        batch_size = elevation_deg.shape[0]
+        camera_embeddings: Float[Tensor, "B 1 4"] = torch.stack(
+            [
+                torch.deg2rad(self.cfg.cond_elevation_deg - elevation_deg),
+                torch.sin(torch.deg2rad(azimuth_deg - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth_deg - self.cfg.cond_azimuth_deg)),
+                camera_distances - self.cfg.cond_camera_distance,
+            ],
+            dim=-1,
+        )[:, None, :]
+
+        image_camera_embeddings = self.pipe.clip_camera_projection(
+            torch.cat(
+                [
+                    self.clip_image_embeddings.repeat(batch_size, 1, 1),
+                    camera_embeddings,
+                ],
+                dim=-1,
+            ).to(self.weights_dtype)
+        )
+
+        return image_camera_embeddings
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        unet: UNet2DConditionModel,
+        latents: Float[Tensor, "..."],
+        t: Int[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        class_labels: Optional[Float[Tensor, "..."]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Float[Tensor, "..."]] = None,
+        mid_block_additional_residual: Optional[Float[Tensor, "..."]] = None,
+        velocity_to_epsilon: bool = False,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        pred = unet(
+            latents.to(unet.dtype),
+            t.to(unet.dtype),
+            encoder_hidden_states=encoder_hidden_states.to(unet.dtype),
+            class_labels=class_labels,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample
+        if velocity_to_epsilon:
+            pred = latents * self.sigmas[t].view(-1, 1, 1, 1) + pred * self.alphas[
+                t
+            ].view(-1, 1, 1, 1)
+        return pred.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_encode(
+        self, vae: AutoencoderKL, imgs: Float[Tensor, "B 3 H W"], mode=False
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        # expect input in [-1, 1]
+        input_dtype = imgs.dtype
+        posterior = vae.encode(imgs.to(vae.dtype)).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample()
+        latents = latents * vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_decode(
+        self, vae: AutoencoderKL, latents: Float[Tensor, "B 4 Hl Wl"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        # output in [0, 1]
+        input_dtype = latents.dtype
+        latents = 1 / vae.config.scaling_factor * latents
+        image = vae.decode(latents.to(vae.dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @contextmanager
+    def disable_unet_class_embedding(self, unet: UNet2DConditionModel):
+        class_embedding = unet.class_embedding
+        try:
+            unet.class_embedding = None
+            yield unet
+        finally:
+            unet.class_embedding = class_embedding
+
+    @contextmanager
+    def set_scheduler(self, pipe: Zero123Pipeline, scheduler_class: Any, **kwargs):
+        scheduler_orig = pipe.scheduler
+        pipe.scheduler = scheduler_class.from_config(scheduler_orig.config, **kwargs)
+        yield pipe
+        pipe.scheduler = scheduler_orig
+
+    def get_eps_pretrain(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        with torch.no_grad():
+            with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                noise_pred = self.forward_unet(
+                    unet,
+                    torch.cat(
+                        [
+                            torch.cat([latents_noisy] * 2, dim=0),
+                            torch.cat(
+                                [
+                                    self.image_latents.repeat(batch_size, 1, 1, 1),
+                                    torch.zeros_like(self.image_latents).repeat(
+                                        batch_size, 1, 1, 1
+                                    ),
+                                ],
+                                dim=0,
+                            ),
+                        ],
+                        dim=1,
+                    ),
+                    torch.cat([t] * 2, dim=0),
+                    encoder_hidden_states=torch.cat(
+                        [
+                            image_camera_embeddings,
+                            torch.zeros_like(image_camera_embeddings),
+                        ],
+                        dim=0,
+                    ),
+                    cross_attention_kwargs={"scale": 0.0}
+                    if self.vsd_share_model
+                    else None,
+                    velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                    == "v_prediction",
+                )
+
+        noise_pred_image, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_image - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def get_eps_phi(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        with torch.no_grad():
+            noise_pred = self.forward_unet(
+                self.pipe_phi.unet,
+                torch.cat(
+                    [
+                        torch.cat([latents_noisy] * 2, dim=0),
+                        torch.cat(
+                            [self.image_latents.repeat(batch_size, 1, 1, 1)] * 2,
+                            dim=0,
+                        ),
+                    ],
+                    dim=1,
+                ),
+                torch.cat([t] * 2, dim=0),
+                encoder_hidden_states=torch.cat([image_camera_embeddings] * 2, dim=0),
+                class_labels=torch.cat(
+                    [
+                        camera_condition.view(batch_size, -1),
+                        torch.zeros_like(camera_condition.view(batch_size, -1)),
+                    ],
+                    dim=0,
+                )
+                if self.cfg.vsd_use_camera_condition
+                else None,
+                cross_attention_kwargs={"scale": 1.0},
+                velocity_to_epsilon=self.pipe_phi.scheduler.config.prediction_type
+                == "v_prediction",
+            )
+
+        noise_pred_camera, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.vsd_guidance_scale_phi * (
+            noise_pred_camera - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def train_phi(
+        self,
+        latents: Float[Tensor, "B 4 Hl Wl"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ):
+        B = latents.shape[0]
+        latents = latents.detach().repeat(
+            self.cfg.vsd_lora_n_timestamp_samples, 1, 1, 1
+        )
+
+        num_train_timesteps = self.pipe_phi.scheduler.config.num_train_timesteps
+        t = torch.randint(
+            int(num_train_timesteps * 0.0),
+            int(num_train_timesteps * 1.0),
+            [B * self.cfg.vsd_lora_n_timestamp_samples],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        noise = torch.randn_like(latents)
+        latents_noisy = self.pipe_phi.scheduler.add_noise(latents, noise, t)
+        if self.pipe_phi.scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif self.pipe_phi.scheduler.prediction_type == "v_prediction":
+            target = self.pipe_phi.scheduler.get_velocity(latents, noise, t)
+        else:
+            raise ValueError(
+                f"Unknown prediction type {self.pipe_phi.scheduler.prediction_type}"
+            )
+
+        if (
+            self.cfg.vsd_use_camera_condition
+            and self.cfg.vsd_lora_cfg_training
+            and random.random() < 0.1
+        ):
+            camera_condition = torch.zeros_like(camera_condition)
+
+        noise_pred = self.forward_unet(
+            self.pipe_phi.unet,
+            torch.cat([latents_noisy, self.image_latents.repeat(B, 1, 1, 1)], dim=1),
+            t,
+            encoder_hidden_states=image_camera_embeddings.repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1, 1
+            ),
+            class_labels=camera_condition.view(B, -1).repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1
+            )
+            if self.cfg.vsd_use_camera_condition
+            else None,
+            cross_attention_kwargs={"scale": 1.0},
+        )
+        return F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+    def forward(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        mvp_mtx: Float[Tensor, "B 4 4"],
+        c2w: Float[Tensor, "B 4 4"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 32 32"]
+        if rgb_as_latents:
+            # treat input rgb as latents
+            # input rgb should be in range [-1, 1]
+            latents = F.interpolate(
+                rgb_BCHW, (32, 32), mode="bilinear", align_corners=False
+            )
+        else:
+            # treat input rgb as rgb
+            # input rgb should be in range [0, 1]
+            rgb_BCHW = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.vae_encode(self.pipe.vae, rgb_BCHW * 2.0 - 1.0)
+
+        # sample timestep
+        # use the same timestep for each batch
+        assert self.min_step is not None and self.max_step is not None
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [1],
+            dtype=torch.long,
+            device=self.device,
+        ).repeat(batch_size)
+
+        # sample noise
+        noise = torch.randn_like(latents)
+        latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+        # image-camera feature condition
+        image_camera_embeddings = self.get_image_camera_embeddings(
+            elevation, azimuth, camera_distances
+        )
+
+        eps_pretrain = self.get_eps_pretrain(
+            latents_noisy,
+            t,
+            image_camera_embeddings,
+            elevation,
+            azimuth,
+            camera_distances,
+        )
+
+        latents_1step_orig = (
+            1
+            / self.alphas[t].view(-1, 1, 1, 1)
+            * (latents_noisy - self.sigmas[t].view(-1, 1, 1, 1) * eps_pretrain)
+        ).detach()
+
+        if self.cfg.guidance_type == "sds":
+            eps_phi = noise
+        elif self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_camera_condition_type == "extrinsics":
+                camera_condition = c2w
+            elif self.cfg.vsd_camera_condition_type == "mvp":
+                camera_condition = mvp_mtx
+            elif self.cfg.vsd_camera_condition_type == "spherical":
+                camera_condition = torch.stack(
+                    [
+                        torch.deg2rad(elevation),
+                        torch.sin(torch.deg2rad(azimuth)),
+                        torch.cos(torch.deg2rad(azimuth)),
+                        camera_distances,
+                    ],
+                    dim=-1,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown camera_condition_type {self.cfg.vsd_camera_condition_type}"
+                )
+            eps_phi = self.get_eps_phi(
+                latents_noisy,
+                t,
+                image_camera_embeddings,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+            loss_train_phi = self.train_phi(
+                latents,
+                image_camera_embeddings,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+        if self.cfg.weighting_strategy == "dreamfusion":
+            w = (1.0 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1.0
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (eps_pretrain - eps_phi)
+
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # reparameterization trick:
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        target = (latents - grad).detach()
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "timesteps": t,
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+            "latents": latents,
+            "latents_1step_orig": latents_1step_orig,
+            "rgb": rgb_BCHW.permute(0, 2, 3, 1),
+            "weights": w,
+            "lambdas": self.lambdas[t],
+        }
+
+        if self.cfg.return_rgb_1step_orig:
+            with torch.no_grad():
+                rgb_1step_orig = self.vae_decode(
+                    self.pipe.vae, latents_1step_orig
+                ).permute(0, 2, 3, 1)
+            guidance_out.update({"rgb_1step_orig": rgb_1step_orig})
+
+        if self.cfg.return_rgb_multistep_orig:
+            with self.set_scheduler(
+                self.pipe,
+                DPMSolverSinglestepScheduler,
+                solver_order=1,
+                num_train_timesteps=int(t[0]),
+            ) as pipe:
+                with torch.cuda.amp.autocast(enabled=False):
+                    latents_multistep_orig = pipe(
+                        num_inference_steps=self.cfg.n_rgb_multistep_orig_steps,
+                        guidance_scale=self.cfg.guidance_scale,
+                        eta=1.0,
+                        latents=latents_noisy.to(pipe.unet.dtype),
+                        image_camera_embeddings=image_camera_embeddings.to(
+                            pipe.unet.dtype
+                        ),
+                        image_latents=self.image_latents.repeat(batch_size, 1, 1, 1).to(
+                            pipe.unet.dtype
+                        ),
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        output_type="latent",
+                    ).images.to(latents.dtype)
+            with torch.no_grad():
+                rgb_multistep_orig = self.vae_decode(
+                    self.pipe.vae, latents_multistep_orig
+                )
+            guidance_out.update(
+                {
+                    "latents_multistep_orig": latents_multistep_orig,
+                    "rgb_multistep_orig": rgb_multistep_orig.permute(0, 2, 3, 1),
+                }
+            )
+
+        if self.cfg.guidance_type == "vsd":
+            guidance_out.update(
+                {
+                    "loss_train_phi": loss_train_phi,
+                }
+            )
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.min_step = int(
+            self.num_train_timesteps * C(self.cfg.min_step_percent, epoch, global_step)
+        )
+        self.max_step = int(
+            self.num_train_timesteps * C(self.cfg.max_step_percent, epoch, global_step)
+        )