float16/float32

app.py

@@ -16,18 +16,20 @@ else:
     device = "cpu"
 model_id = "stabilityai/stable-diffusion-2-1-base"
 
+to = torch.float16 if self.device == 'cuda' else torch.float32
+
 # components for the Preprocessor
 scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
 vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", revision="fp16",
-                                                 torch_dtype=torch.float16).to(device)
+                                                 torch_dtype=to).to(device)
 tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
 text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision="fp16",
-                                                  torch_dtype=torch.float16).to(device)
+                                                  torch_dtype=to).to(device)
 unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", revision="fp16",
-                                           torch_dtype=torch.float16).to(device)
+                                           torch_dtype=to).to(device)
 
 # pipe for TokenFlow
-tokenflow_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
+tokenflow_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=to).to(device)
 if device == "cuda":
     tokenflow_pipe.enable_xformers_memory_efficient_attention()
 

preprocess_utils.py

@@ -29,6 +29,7 @@ class Preprocess(nn.Module):
         super().__init__()
 
         self.device = device
+        self.to = torch.float16 if self.device == 'cuda' else torch.float32
         self.sd_version = opt["sd_version"]
         self.use_depth = False
         self.config = opt
@@ -73,9 +74,9 @@ class Preprocess(nn.Module):
         
         if self.sd_version == 'ControlNet':
             from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
-            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device)
+            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=self.to).to(self.device)
             control_pipe = StableDiffusionControlNetPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=self.to
             ).to(self.device)
             self.unet = control_pipe.unet
             self.controlnet = control_pipe.controlnet
@@ -124,7 +125,7 @@ class Preprocess(nn.Module):
             depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
             depth_maps.append(depth_map)
 
-        return torch.cat(depth_maps).to(self.device).to(torch.float16)
+        return torch.cat(depth_maps).to(self.device).to(self.to)
     
     @torch.no_grad()
     def get_canny_cond(self):
@@ -139,7 +140,7 @@ class Preprocess(nn.Module):
             image = np.concatenate([image, image, image], axis=2)
             image = torch.from_numpy((image.astype(np.float32) / 255.0))
             canny_cond.append(image)
-        canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(torch.float16)
+        canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(self.to)
         return canny_cond
     
     def controlnet_pred(self, latent_model_input, t, text_embed_input, controlnet_cond):
@@ -211,9 +212,9 @@ class Preprocess(nn.Module):
                 frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
         else:
             frames = self.config["frames"][:n_frames]
-        frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
+        frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(self.to).to(self.device)
         # encode to latents
-        latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
+        latents = self.encode_imgs(frames, deterministic=True).to(self.to).to(self.device)
         print("frames", frames.shape)
         print("latents", latents.shape)
         

tokenflow_pnp.py

@@ -1,3 +1,4 @@
+import torch
 import glob
 import os
 import numpy as np
@@ -21,6 +22,12 @@ logging.set_verbosity_error()
 
 VAE_BATCH_SIZE = 10
 
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
 
 class TokenFlow(nn.Module):
     def __init__(self, config, 
@@ -31,6 +38,7 @@ class TokenFlow(nn.Module):
         super().__init__()
         self.config = config
         self.device = config["device"]
+        self.to = torch.float16 if self.device == 'cuda' else torch.float32
         
         sd_version = config["sd_version"]
         self.sd_version = sd_version
@@ -109,7 +117,7 @@ class TokenFlow(nn.Module):
             depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
             depth_maps.append(depth_map)
 
-        return torch.cat(depth_maps).to(torch.float16).to(self.device)
+        return torch.cat(depth_maps).to(self.to).to(self.device)
     
     def get_pnp_inversion_prompt(self):
         inv_prompts_path = os.path.join(str(Path(self.latents_path).parent), 'inversion_prompt.txt')
@@ -197,16 +205,16 @@ class TokenFlow(nn.Module):
             frames = [Image.open(paths[idx]).convert('RGB') for idx in range(self.config["n_frames"])]
             if frames[0].size[0] == frames[0].size[1]:
                 frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
-            frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
+            frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(self.to).to(self.device)
             save_video(frames, f'{self.config["output_path"]}/input_fps10.mp4', fps=10)
             save_video(frames, f'{self.config["output_path"]}/input_fps20.mp4', fps=20)
             save_video(frames, f'{self.config["output_path"]}/input_fps30.mp4', fps=30)
         else:
             frames = self.frames
         # encode to latents
-        latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
+        latents = self.encode_imgs(frames, deterministic=True).to(self.to).to(self.device)
         # get noise
-        eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(torch.float16).to(self.device)
+        eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(self.to).to(self.device)
         if not read_from_files:
             return None, frames, latents, eps
         return paths, frames, latents, eps
@@ -267,7 +275,7 @@ class TokenFlow(nn.Module):
         denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
         return denoised_latent
     
-    @torch.autocast(dtype=torch.float16, device_type='cuda')
+    @torch.autocast(dtype=self.to, device_type=device)
     def batched_denoise_step(self, x, t, indices):
         batch_size = self.config["batch_size"]
         denoised_latents = []