wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-10 13:32:36 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Make torch.compile do something properly and work once eigenops is upgraded.

Browse Source
This commit is contained in:
Simon Lui 2023-08-24 23:37:37 -07:00
parent cf4d7c485f
commit 0e2a834c86
3 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/cli_args.py
View File

@ -61,7 +61,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize when loading models with Intel GPUs.")
parser.add_argument("--disable-torch-compile", action="store_true", help="Disables torch.compile for loading models.")
parser.add_argument("--torch-compile-fullgraph", action="store_true", default=True, help="torch.compile argument for if the model should be compiled into a single graph.")
parser.add_argument("--torch-compile-fullgraph", action="store_true", default=False, help="torch.compile argument for if the model should be compiled into a single graph.")
parser.add_argument("--torch-compile-backend", type=str, default="inductor", help="torch.compile argument for what backend to use. See Pytorch documentation for available backends to choose from.")
parser.add_argument("--torch-compile-mode", type=str, default="default", help="torch.compile argument for what compile mode to use. Options include 'default', 'reduce-overhead', or 'max-autotune'.")

3
comfy/model_management.py
View File

@ -273,9 +273,6 @@ class LoadedModel:
if xpu_available and not args.disable_ipex_optimize:
self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
if not args.disable_torch_compile:
self.real_model = torch.compile(self.real_model, fullgraph=args.torch_compile_fullgraph, backend=args.torch_compile_backend, mode=args.torch_compile_mode)
return self.real_model
def model_unload(self):

5
comfy/sample.py
View File

@ -1,3 +1,4 @@
from comfy.cli_args import args
import torch
import comfy.model_management
import comfy.samplers
@ -81,13 +82,15 @@ def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative
comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]))
real_model = model.model
if not args.disable_torch_compile:
real_model.diffusion_model = torch.compile(real_model.diffusion_model, fullgraph=args.torch_compile_fullgraph, backend=args.torch_compile_backend, mode=args.torch_compile_mode)
noise = noise.to(device)
latent_image = latent_image.to(device)
positive_copy = broadcast_cond(positive, noise.shape[0], device)
negative_copy = broadcast_cond(negative, noise.shape[0], device)
sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)