|
|
|
@ -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,
|
|
|
|
|