From 2da73b7073dc520ee480dee8ff911b9aa83ff70a Mon Sep 17 00:00:00 2001
From: Simon Lui <502929+simonlui@users.noreply.github.com>
Date: Sat, 2 Sep 2023 20:07:52 -0700
Subject: [PATCH] Revert changes in comfy/ldm/modules/diffusionmodules/util.py,
 which is unused.

---
 comfy/ldm/modules/diffusionmodules/util.py | 24 +++++++---------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/comfy/ldm/modules/diffusionmodules/util.py b/comfy/ldm/modules/diffusionmodules/util.py
index 9d07d93..d890c80 100644
--- a/comfy/ldm/modules/diffusionmodules/util.py
+++ b/comfy/ldm/modules/diffusionmodules/util.py
@@ -15,7 +15,6 @@ import torch.nn as nn
 import numpy as np
 from einops import repeat
 
-from comfy import model_management
 from comfy.ldm.util import instantiate_from_config
 import comfy.ops
 
@@ -140,22 +139,13 @@ class CheckpointFunction(torch.autograd.Function):
     @staticmethod
     def backward(ctx, *output_grads):
         ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
-        if model_management.is_nvidia():
-            with torch.enable_grad(), \
-                    torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
-                # Fixes a bug where the first op in run_function modifies the
-                # Tensor storage in place, which is not allowed for detach()'d
-                # Tensors.
-                shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
-                output_tensors = ctx.run_function(*shallow_copies)
-        elif model_management.is_intel_xpu():
-            with torch.enable_grad(), \
-                    torch.xpu.amp.autocast(**ctx.gpu_autocast_kwargs):
-                # Fixes a bug where the first op in run_function modifies the
-                # Tensor storage in place, which is not allowed for detach()'d
-                # Tensors.
-                shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
-                output_tensors = ctx.run_function(*shallow_copies)
+        with torch.enable_grad(), \
+                torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
         input_grads = torch.autograd.grad(
             output_tensors,
             ctx.input_tensors + ctx.input_params,