Back to float32

app.py

@@ -30,23 +30,20 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # }
 # """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
+if device == "cuda":
     torch.cuda.max_memory_allocated(device=device)
-else:
-    power_device = "CPU"
 
 scheduler_class = MyEulerAncestralDiscreteScheduler
 
-pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)  # .to('cpu')
-pipe_inference = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16",
-                                                            use_safetensors=True)  # .to('cpu')
+pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)  # .to('cpu')
+pipe_inference = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)  # .to('cpu')
 pipe_inference.scheduler = scheduler_class.from_config(pipe_inference.scheduler.config)
 pipe_inversion.scheduler = scheduler_class.from_config(pipe_inversion.scheduler.config)
 pipe_inversion.scheduler_inference = scheduler_class.from_config(pipe_inference.scheduler.config)
 
-pipe_inference.enable_xformers_memory_efficient_attention()
-pipe_inversion.enable_xformers_memory_efficient_attention()
+if device == "cuda":
+    pipe_inference.enable_xformers_memory_efficient_attention()
+    pipe_inversion.enable_xformers_memory_efficient_attention()
 
 
 

src/editor.py

@@ -44,7 +44,7 @@ class ImageEditorDemo:
         img_size = (512,512)
         VQAE_SCALE = 8
         latents_size = (1, 4, img_size[0] // VQAE_SCALE, img_size[1] // VQAE_SCALE)
-        noise = [randn_tensor(latents_size, dtype=torch.float16, device=torch.device(device), generator=g_cpu) for i
+        noise = [randn_tensor(latents_size, dtype=torch.float32, device=torch.device(device), generator=g_cpu) for i
                  in range(cfg.num_inversion_steps)]
         pipe_inversion.scheduler.set_noise_list(noise)
         pipe_inversion.scheduler_inference.set_noise_list(noise)

src/euler_scheduler.py

@@ -120,7 +120,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
         sigma = self.sigmas[self.step_index]
 
         # Upcast to avoid precision issues when computing prev_sample
-        # sample = sample.to(torch.float32)
+        sample = sample.to(torch.float32)
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
         if self.config.prediction_type == "epsilon":
@@ -226,7 +226,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
         sigma = self.sigmas[self.step_index]
 
         # Upcast to avoid precision issues when computing prev_sample
-        # sample = sample.to(torch.float32)
+        sample = sample.to(torch.float32)
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
         if self.config.prediction_type == "epsilon":
@@ -342,7 +342,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
         sigma = self.sigmas[self.step_index]
 
         # Upcast to avoid precision issues when computing prev_sample
-        # sample = sample.to(torch.float32)
+        sample = sample.to(torch.float32)
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
         if self.config.prediction_type == "epsilon":
@@ -394,7 +394,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
 
     def get_all_sigmas(self) -> torch.FloatTensor:
         sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
-        sigmas = np.concatenate([sigmas[::-1], [0.0]])#.astype(np.float32)
+        sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
         return torch.from_numpy(sigmas)
 
     def add_noise_off_schedule(
@@ -408,7 +408,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
         sigmas = sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
         if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
             # mps does not support float64
-            timesteps = timesteps.to(original_samples)#.device, dtype=torch.float32)
+            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
         else:
             timesteps = timesteps.to(original_samples.device)
 
@@ -539,7 +539,7 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
     #     sigma = self.sigmas[self.step_index]
 
     #     # Upcast to avoid precision issues when computing prev_sample
-    #     sample = sample.to(torch.float32)
+        sample = sample.to(torch.float32)
 
     #     # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
     #     if self.config.prediction_type == "epsilon":