add deepfloyd model (#1993)

* add deepfloyd model

* black

* lint

* linting

* cast sds gradient to tensor

* change generator type to optional

* lint

* Type fixes

* add typing to sds

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>
Co-authored-by: Matt Tancik <mattt4001@gmail.com>

nerfstudio/generative/deepfloyd.py

@@ -0,0 +1,236 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+from pathlib import Path
+from typing import List, Optional, Union, cast
+
+import torch
+import torch.nn.functional as F
+import tyro
+from diffusers import IFPipeline
+from diffusers.pipelines.deepfloyd_if import IFPipelineOutput
+from jaxtyping import Float
+from PIL import Image
+from torch import FloatTensor, Generator, Tensor, nn
+from torch.cuda.amp.grad_scaler import GradScaler
+
+from transformers import T5EncoderModel
+
+from nerfstudio.generative.utils import _SDSGradient
+
+IMG_DIM = 64
+
+
+class DeepFloyd(nn.Module):
+    """DeepFloyd diffusion model
+    Args:
+        device: device to use
+    """
+
+    def __init__(self, device: Union[torch.device, str]):
+        super().__init__()
+        self.device = device
+
+        self.text_encoder = T5EncoderModel.from_pretrained(
+            "DeepFloyd/IF-I-L-v1.0",
+            subfolder="text_encoder",
+            load_in_8bit=True,
+            variant="8bit",
+            device_map="auto",
+        )
+
+        self.pipe = IFPipeline.from_pretrained(
+            "DeepFloyd/IF-I-L-v1.0",
+            text_encoder=self.text_encoder,
+            safety_checker=None,
+            watermarker=None,
+            feature_extractor=None,
+            requires_safety_checker=False,
+            variant="fp16",
+            torch_dtype=torch.float16,
+        )
+        assert isinstance(self.pipe, IFPipeline)
+        self.pipe = self.pipe.to(self.device)
+
+        self.pipe.enable_attention_slicing(1)
+
+        self.unet = self.pipe.unet
+        self.unet.to(memory_format=torch.channels_last)  # type: ignore
+        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.min_step = int(self.num_train_timesteps * 0.02)
+        self.max_step = int(self.num_train_timesteps * 0.98)
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(self.device)
+
+    def delete_text_encoder(self):
+        """Delete text encoder from pipeline. T5 text encoder uses a lot of memory."""
+        del self.text_encoder
+        del self.pipe
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        self.pipe = IFPipeline.from_pretrained(
+            "DeepFloyd/IF-I-L-v1.0",
+            text_encoder=None,
+            safety_checker=None,
+            watermarker=None,
+            feature_extractor=None,
+            requires_safety_checker=False,
+            variant="fp16",
+            torch_dtype=torch.float16,
+        )
+        assert isinstance(self.pipe, IFPipeline)
+        self.pipe = self.pipe.to(self.device)
+
+        self.pipe.enable_attention_slicing(1)
+
+        self.unet = self.pipe.unet
+        self.unet.to(memory_format=torch.channels_last)  # type: ignore
+
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        self.scheduler = self.pipe.scheduler
+
+    def get_text_embeds(
+        self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
+    ) -> Float[Tensor, "2 max_length embed_dim"]:
+        """Get text embeddings for prompt and negative prompt
+        Args:
+            prompt: Prompt text
+            negative_prompt: Negative prompt text
+        Returns:
+            Text embeddings
+        """
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+        assert isinstance(self.pipe, IFPipeline)
+        with torch.no_grad():
+            prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompt, negative_prompt=negative_prompt)
+
+        assert isinstance(negative_embeds, FloatTensor)
+        assert isinstance(prompt_embeds, FloatTensor)
+        return torch.cat([negative_embeds, prompt_embeds])
+
+    def sds_loss(
+        self,
+        text_embeddings: Float[Tensor, "N max_length embed_dim"],
+        image: Float[Tensor, "BS 3 H W"],
+        guidance_scale: float = 100.0,
+        grad_scaler: Optional[GradScaler] = None,
+    ) -> torch.Tensor:
+        """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
+        Args:
+            text_embeddings: Text embeddings
+            image: Rendered image
+            guidance_scale: How much to weigh the guidance
+            grad_scaler: Grad scaler
+        Returns:
+            The loss
+        """
+        with torch.autocast(device_type="cuda", enabled=False):
+            image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear", align_corners=False)
+            t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
+
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(image)
+                image_noisy = self.scheduler.add_noise(image, noise, t)  # type: ignore
+                # pred noise
+                image_model_input = torch.cat((image_noisy,) * 2)
+                noise_pred = self.unet(image_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_text = 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 + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # w(t), sigma_t^2
+            w = 1 - self.alphas[t]
+
+            grad = w * (noise_pred - noise)
+            grad = torch.nan_to_num(grad)
+
+            if grad_scaler is not None:
+                loss = cast(Tensor, _SDSGradient.apply(grad_scaler.scale(image), grad))
+            else:
+                loss = cast(Tensor, _SDSGradient.apply(image, grad))
+
+        return loss
+
+    def prompt_to_img(
+        self,
+        prompts: Union[str, List[str]],
+        negative_prompts: Union[str, List[str]] = "",
+        generator: Optional[Generator] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+    ) -> Image.Image:
+        """Generate an image from a prompt.
+        Args:
+            prompts: The prompt to generate an image from.
+            negative_prompts: The negative prompt to generate an image from.
+            generator: Random seed.
+            num_inference_steps: The number of inference steps to perform.
+            guidance_scale: The scale of the guidance.
+            latents: The latents to start from, defaults to random.
+        Returns:
+            The generated image.
+        """
+
+        prompts = [prompts] if isinstance(prompts, str) else prompts
+        negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts
+
+        assert isinstance(self.pipe, IFPipeline)
+        prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompts, negative_prompt=negative_prompts)
+        model_output = self.pipe(
+            prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator
+        )
+        assert isinstance(model_output, IFPipelineOutput)
+        image = model_output.images[0]
+
+        return image
+
+
+def generate_image(
+    prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_deepfloyd.png")
+):
+    """Generate an image from a prompt using DeepFloyd IF.
+    Args:
+        prompt: The prompt to use.
+        negative: The negative prompt to use.
+        seed: The random seed to use.
+        steps: The number of steps to use.
+        save_path: The path to save the image to.
+    """
+    generator = torch.manual_seed(seed)
+    cuda_device = torch.device("cuda")
+    with torch.no_grad():
+        df = DeepFloyd(cuda_device)
+        img = df.prompt_to_img(prompt, negative, generator, steps)
+        img.save(save_path)
+
+
+if __name__ == "__main__":
+    tyro.cli(generate_image)

nerfstudio/generative/stable_diffusion.py

@@ -28,11 +28,13 @@ import torch.nn.functional as F
 import tyro
 from jaxtyping import Float
 from torch import Tensor, nn
-from torch.cuda.amp import custom_bwd, custom_fwd
 from torch.cuda.amp.grad_scaler import GradScaler
 
 from nerfstudio.utils.rich_utils import CONSOLE
 
+from nerfstudio.generative.utils import _SDSGradient
+
+
 try:
     from diffusers import PNDMScheduler, StableDiffusionPipeline
     from transformers import logging
@@ -60,26 +62,6 @@ class UNet2DConditionOutput:
     sample: torch.FloatTensor
 
 
-class _SDSGradient(torch.autograd.Function):
-    """Custom gradient function for SDS loss. Since it is already computed, we can just return it."""
-
-    @staticmethod
-    @custom_fwd
-    def forward(ctx, input_tensor, gt_grad):
-        del input_tensor
-        ctx.save_for_backward(gt_grad)
-        # Return magniture of gradient, not the actual loss.
-        return torch.mean(gt_grad**2) ** 0.5
-
-    @staticmethod
-    @custom_bwd
-    def backward(ctx, grad):
-        del grad
-        (gt_grad,) = ctx.saved_tensors
-        batch_size = len(gt_grad)
-        return gt_grad / batch_size, None
-
-
 class StableDiffusion(nn.Module):
     """Stable Diffusion implementation
     Args:
@@ -224,7 +206,8 @@ class StableDiffusion(nn.Module):
 
         if latents is None:
             latents = torch.randn(
-                (text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device
+                (text_embeddings.shape[0] // 2, self.unet.config.in_channels, height // 8, width // 8),
+                device=self.device,
             )
 
         self.scheduler.set_timesteps(num_inference_steps)  # type: ignore

nerfstudio/generative/utils.py

@@ -0,0 +1,39 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+"""Gradient helper functions for SDS loss"""
+
+
+class _SDSGradient(torch.autograd.Function):
+    """Custom gradient function for SDS loss. Since it is already computed, we can just return it."""
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, input_tensor, gt_grad):
+        del input_tensor
+        ctx.save_for_backward(gt_grad)
+        # Return magniture of gradient, not the actual loss.
+        return torch.mean(gt_grad**2) ** 0.5
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        del grad
+        (gt_grad,) = ctx.saved_tensors
+        batch_size = len(gt_grad)
+        return gt_grad / batch_size, None