Merge branch 'comfyanonymous:master' into feature/settings

comfy/cli_args.py

@@ -40,6 +40,10 @@ parser.add_argument("--extra-model-paths-config", type=str, default=None, metava
 parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
 parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
 parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
+cm_group = parser.add_mutually_exclusive_group()
+cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
+cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Enable cudaMallocAsync.")
+
 parser.add_argument("--dont-upcast-attention", action="store_true", help="Disable upcasting of attention. Can boost speed but increase the chances of black images.")
 
 fp_group = parser.add_mutually_exclusive_group()

comfy/k_diffusion/external.py

@@ -92,8 +92,8 @@ class DiscreteSchedule(nn.Module):
 
     def predict_eps_discrete_timestep(self, input, t, **kwargs):
         sigma = self.t_to_sigma(t.round())
-        input = input * ((sigma ** 2 + 1.0) ** 0.5)
-        return  (input - self(input, sigma, **kwargs)) / sigma
+        input = input * ((utils.append_dims(sigma, input.ndim) ** 2 + 1.0) ** 0.5)
+        return  (input - self(input, sigma, **kwargs)) / utils.append_dims(sigma, input.ndim)
 
 class DiscreteEpsDDPMDenoiser(DiscreteSchedule):
     """A wrapper for discrete schedule DDPM models that output eps (the predicted

comfy/k_diffusion/sampling.py

@@ -131,9 +131,9 @@ def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None,
     s_in = x.new_ones([x.shape[0]])
     for i in trange(len(sigmas) - 1, disable=disable):
         gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = torch.randn_like(x) * s_noise
         sigma_hat = sigmas[i] * (gamma + 1)
         if gamma > 0:
+            eps = torch.randn_like(x) * s_noise
             x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
         denoised = model(x, sigma_hat * s_in, **extra_args)
         d = to_d(x, sigma_hat, denoised)
@@ -172,9 +172,9 @@ def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None,
     s_in = x.new_ones([x.shape[0]])
     for i in trange(len(sigmas) - 1, disable=disable):
         gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = torch.randn_like(x) * s_noise
         sigma_hat = sigmas[i] * (gamma + 1)
         if gamma > 0:
+            eps = torch.randn_like(x) * s_noise
             x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
         denoised = model(x, sigma_hat * s_in, **extra_args)
         d = to_d(x, sigma_hat, denoised)
@@ -201,9 +201,9 @@ def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None,
     s_in = x.new_ones([x.shape[0]])
     for i in trange(len(sigmas) - 1, disable=disable):
         gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = torch.randn_like(x) * s_noise
         sigma_hat = sigmas[i] * (gamma + 1)
         if gamma > 0:
+            eps = torch.randn_like(x) * s_noise
             x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
         denoised = model(x, sigma_hat * s_in, **extra_args)
         d = to_d(x, sigma_hat, denoised)

comfy/model_management.py

@@ -204,7 +204,11 @@ print(f"Set vram state to: {vram_state.name}")
 def get_torch_device_name(device):
     if hasattr(device, 'type'):
         if device.type == "cuda":
-            return "{} {}".format(device, torch.cuda.get_device_name(device))
+            try:
+                allocator_backend = torch.cuda.get_allocator_backend()
+            except:
+                allocator_backend = ""
+            return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
         else:
             return "{}".format(device.type)
     else:

main.py

@@ -51,7 +51,6 @@ import threading
 import gc
 
 from comfy.cli_args import args
-import comfy.utils
 
 if os.name == "nt":
     import logging
@@ -62,7 +61,32 @@ if __name__ == "__main__":
         os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
         print("Set cuda device to:", args.cuda_device)
 
+    if not args.cuda_malloc:
+        try: #if there's a better way to check the torch version without importing it let me know
+            version = ""
+            torch_spec = importlib.util.find_spec("torch")
+            for folder in torch_spec.submodule_search_locations:
+                ver_file = os.path.join(folder, "version.py")
+                if os.path.isfile(ver_file):
+                    spec = importlib.util.spec_from_file_location("torch_version_import", ver_file)
+                    module = importlib.util.module_from_spec(spec)
+                    spec.loader.exec_module(module)
+                    version = module.__version__
+            if int(version[0]) >= 2: #enable by default for torch version 2.0 and up
+                args.cuda_malloc = True
+        except:
+            pass
 
+    if args.cuda_malloc and not args.disable_cuda_malloc:
+        env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)
+        if env_var is None:
+            env_var = "backend:cudaMallocAsync"
+        else:
+            env_var += ",backend:cudaMallocAsync"
+
+        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var
+
+import comfy.utils
 import yaml
 
 import execution