add deepfloyd option first try

nerfstudio/generative/deepfloyd.py

@@ -0,0 +1,162 @@
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Union
+
+import appdirs
+import mediapy
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tyro
+from rich.console import Console
+from torch import nn
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.cuda.amp.grad_scaler import GradScaler
+from torchtyping import TensorType
+
+
+
+from diffusers import IFPipeline
+from transformers import T5EncoderModel, T5Tokenizer
+
+class DeepFloyd(nn.Module):
+
+    # load in model
+    def __init__(self, device: Union[torch.device, str]):
+        super().__init__()
+        self.device = device
+
+        self.text_encoder = T5EncoderModel.from_pretrained(
+            "DeepFloyd/IF-I-M-v1.0",
+            subfolder="text_encoder",
+            load_in_8bit=True,
+            variant="8bit",
+            device_map="auto",
+        )
+
+        self.pipe = IFPipeline.from_pretrained(
+            "DeepFloyd/IF-I-M-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,
+        ).to(self.device)
+        
+        self.pipe.enable_attention_slicing(1)
+        self.pipe.unet.to(memory_format=torch.channels_last)
+
+        self.unet = self.pipe.unet
+        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 get_text_embeds(
+        self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
+    ) -> TensorType[2, "max_length", "embed_dim"]:        
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+        prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompt, negative_prompt=negative_prompt)
+
+        return torch.cat([prompt_embeds, negative_embeds])
+
+    def sds_loss(
+        self,
+        text_embeddings, 
+        image, 
+        guidance_scale, 
+        grad_scaler,
+    ) -> torch.Tensor:
+        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_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 + 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:
+                latents = grad_scaler.scale(latents)
+            loss = _SDSGradient.apply(latents, grad)
+
+        return loss
+
+    def prompt_to_img(
+        self,
+        prompts: Union[str, List[str]],
+        negative_prompts: Union[str, List[str]] = "",
+        generator=None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+    ) -> np.ndarray:
+        """Generate an images from a prompts.
+        Args:
+            prompts: The prompt to generate an image from.
+            negative_prompts: The negative prompt to generate an image from.
+            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
+        # text_embeddings = self.get_text_embeds(prompts, negative_prompts)
+        prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompts, negative_prompt=negative_prompts)
+        image = self.pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator).images
+
+        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)
+        imgs = df.prompt_to_img(prompt, negative, generator, steps)
+        imgs[0].save(save_path)
+        # mediapy.write_image(str(save_path), imgs[0])
+
+if __name__ == "__main__":
+    tyro.cli(generate_image)

nerfstudio/models/dreamfusion.py

@@ -36,6 +36,7 @@ from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
 from nerfstudio.generative.stable_diffusion import StableDiffusion
+from nerfstudio.generative.deepfloyd import DeepFloyd
 from nerfstudio.generative.stable_diffusion_utils import PositionalTextEmbeddings
 from nerfstudio.model_components.losses import (
     MSELoss,
@@ -198,14 +199,16 @@ class DreamFusionModel(Model):
         """Set the fields and modules"""
         super().populate_modules()
 
-        self.sd = StableDiffusion(self.sd_device, version=self.sd_version)
+        # self.sd = StableDiffusion(self.sd_device, version=self.sd_version)
+        self.df = DeepFloyd(self.sd_device)
         self.text_embeddings = PositionalTextEmbeddings(
             base_prompt=self.cur_prompt,
             top_prompt=self.cur_prompt + self.top_prompt,
             side_prompt=self.cur_prompt + self.side_prompt,
             back_prompt=self.cur_prompt + self.back_prompt,
             front_prompt=self.cur_prompt + self.front_prompt,
-            stable_diffusion=self.sd,
+            # stable_diffusion=self.sd,
+            stable_diffusion=self.df,
             positional_prompting=self.positional_prompting,
         )
 
@@ -480,11 +483,18 @@ class DreamFusionModel(Model):
             .permute(0, 3, 1, 2)
         )
 
-        sds_loss = self.sd.sds_loss(
+        # sds_loss = self.sd.sds_loss(
+        #     text_embedding.to(self.sd_device),
+        #     train_output.to(self.sd_device),
+        #     guidance_scale=int(self.guidance_scale),
+        #     grad_scaler=self.grad_scaler,
+        # )
+
+        sds_loss = self.df.sds_loss(
             text_embedding.to(self.sd_device),
             train_output.to(self.sd_device),
             guidance_scale=int(self.guidance_scale),
-            grad_scaler=self.grad_scaler,
+            grad_scaler=self.grad_scaler,        
         )
 
         loss_dict["sds_loss"] = sds_loss.to(self.device)