Compiling modules on Linux

[hlky]

This issue can serve as a discussion point for any issues related to compiling modules on Linux

My setup is as follows:

I use WSL, with Ubuntu 20.04

wsl --install Ubuntu-20.04
ubuntu2004.exe

CUDA installation is specific to WSL, for other distros see cuda-downloads.
Earliest known working CUDA toolkit version is 11.6, I'd recommend using the latest

wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.2.1/local_installers/cuda-repo-wsl-ubuntu-12-2-local_12.2.1-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-2-local_12.2.1-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-12-2-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

Ubuntu 20.04 includes Python 3.8.10, minimum Python version is 3.8, and I use PyTorch nightly CUDA 12.1, PyTorch version working at time of AITemplate release was 1.12.1, so anything newer than that should work
The following instructions use the build of AITemplate mentioned here, it includes the pending PRs, it is set up for CMake so that is disabled here, CUDA path is set so nvcc is found, a venv is used

sudo apt install python3-pip python3-venv unzip
wget https://github.com/FizzleDorf/AIT/files/12283714/aitemplate.zip
unzip aitemplate.zip
python3 -m venv venv
mv aitemplate/ venv/lib/python3.8/site-packages/
venv/bin/python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
venv/bin/python -m pip install click transformers diffusers accelerate
export CUDA_HOME=/usr/local/cuda
export PATH=${CUDA_HOME}/bin:${PATH}
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
export AIT_USE_CMAKE_COMPILATION=0

To note, using that build is not necessary, and it will eventually become outdated, Installation from source is like

git clone --recursive https://github.com/facebookincubator/AITemplate
cd python
python setup.py bdist_wheel
pip install dist/*.whl

Until the pending PRs mentioned are merged they will need to be merged to a local branch before installation, or the changes applied to the installed package.

Here's some SDXL UNet modules I compiled earlier, covering 1024, 1536, 2048 max resolution, batch size 1, sm80
unet_xl_linux.zip

This is actually a 7z archive rather than .zip due to GitHub's file extension constraints.

cc @FizzleDorf These UNet modules are ready uploaded to Huggingface and added to modules.json, existing Linux VAE modules will work for SDXL with the fp16 fixed version

[CyberTimon]

Error occurred when executing KSampler:

/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/modules/b5caabe98aeb69bada9d1566c897aed66a84d4fb21f31482160d7ef9987f04fd.so: file too short

I get this issue when I'm trying to use the latest branch with SDXL. Any ideas?

[hlky]

@CyberTimon
As the error suggests there is an issue with the file. Upon further inspection of modules.json it appears some urls including the one for that module also need adjusting.
To fix this locally, first delete the file from the modules folder, change any occurrences of blob in modules.json to resolve, and try again, the module should then download correctly.

[CyberTimon]

Thank you this fixed this issue. Now hopefully the last issue:

Error occurred when executing KSampler:

Input format not supported by decoder

File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 151, in recursive_execute
output_data, output_ui = get_output_data(obj, input_data_all)
File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 81, in get_output_data
return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 74, in map_node_over_list
results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
File "/home/cybertimon/Repositories/ComfyUI/nodes.py", line 1206, in sample
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 170, in common_ksampler
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 254, in sample
AITemplate.unet[module['sha256']] = AITemplate.loader.load_module(module['sha256'], module['url'])
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/load.py", line 68, in load_module
self.download_module(sha256, url)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/load.py", line 55, in download_module
g.write(f.read())
File "/usr/lib/python3.10/lzma.py", line 200, in read
return self._buffer.read(size)
File "/usr/lib/python3.10/_compression.py", line 118, in readall
while data := self.read(sys.maxsize):
File "/usr/lib/python3.10/_compression.py", line 103, in read
data = self._decompressor.decompress(rawblock, size)

Do you know how I should fix this? And the filename of the downloaded unet is still very random

[CyberTimon]

After furter investigation I noticed that when I first load up comfy with the node it downloads the aitemplate unet and it outputs the first issue (file too short), even after changing blob to resolve. The second time I execute the prompt I get the above input format error.

[hlky]

Ah, yes the the script intends them to be compressed for the download. The compressed versions will be uploaded later.

Instead, download the module yourself, here is the direct link and rename it to the same as the existing file b5caabe98aeb69bada9d1566c897aed66a84d4fb21f31482160d7ef9987f04fd.so. The name may appear random, it is the sha256 of the file, you can see it displayed here on the huggingface page, the same value is in modules.json, in the usual case of compressed modules it is the sha256 not the compressed_sha256 that determines the filename where the module is loaded from. The sha256 is intended to verify the downloaded file is correct, although the script does not actually check I decided to use it for filename and include in the json in case module verification was implemented.

[CyberTimon]

Ahhh it really seems impossible to use AiTemplate on linux on my system. I did exactly what you said. I now have a single b5caabe98aeb69bada9d1566c897aed66a84d4fb21f31482160d7ef9987f04fd.so file in the modules directory but now I get this:

Error occurred when executing KSampler:

Got unexpected input: input0

File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 151, in recursive_execute
output_data, output_ui = get_output_data(obj, input_data_all)
File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 81, in get_output_data
return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
File "/home/cybertimon/Repositories/ComfyUI/execution.py", line 74, in map_node_over_list
results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
File "/home/cybertimon/Repositories/ComfyUI/nodes.py", line 1206, in sample
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 170, in common_ksampler
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 304, in sample
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)
File "/home/cybertimon/Repositories/ComfyUI/comfy/samplers.py", line 716, in sample
samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
File "/home/cybertimon/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/cybertimon/Repositories/ComfyUI/comfy/k_diffusion/sampling.py", line 137, in sample_euler
denoised = model(x, sigma_hat * s_in, **extra_args)
File "/home/cybertimon/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/cybertimon/Repositories/ComfyUI/comfy/samplers.py", line 319, in forward
out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, cond_concat=cond_concat, model_options=model_options, seed=seed)
File "/home/cybertimon/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/cybertimon/Repositories/ComfyUI/comfy/k_diffusion/external.py", line 125, in forward
eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
File "/home/cybertimon/Repositories/ComfyUI/comfy/k_diffusion/external.py", line 151, in get_eps
return self.inner_model.apply_model(*args, **kwargs)
File "/home/cybertimon/Repositories/ComfyUI/comfy/samplers.py", line 307, in apply_model
out = sampling_function(self.inner_model.apply_model, x, timestep, uncond, cond, cond_scale, cond_concat, model_options=model_options, seed=seed)
File "/home/cybertimon/Repositories/ComfyUI/comfy/samplers.py", line 285, in sampling_function
cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options)
File "/home/cybertimon/Repositories/ComfyUI/comfy/samplers.py", line 262, in calc_cond_uncond_batch
output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/inference.py", line 43, in apply_model
return unet_inference(
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/inference.py", line 98, in unet_inference
exe_module.run_with_tensors(inputs, ys, graph_mode=False)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 565, in run_with_tensors
outputs_ait = self.run(
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 468, in run
return self._run_impl(
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 397, in _run_impl
inputs = self._dict_to_ordered_list(inputs, is_inputs=True)
File "/home/cybertimon/Repositories/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 364, in _dict_to_ordered_list
raise ValueError(

[hlky]

See #15

[CyberTimon]

Thank you. I renamed the suggested things and now I get this:

Using b5caabe98aeb69bada9d1566c897aed66a84d4fb21f31482160d7ef9987f04fd
  0%|                                                                                                                                                                                | 0/20 [00:00<?, ?it/s]

Got cutlass error: Error InternalError: Got cutlass error: Error Internal

I think it's the best to wait until you know a real fix or a PR appears.

[David-337]

Same as @CyberTimon, here's the full error:

  0%|                                                                       | 0/35 [00:00<?, ?it/s]Got cutlass error: Error InternalError: Got cutlass error: Error Internal
  0%|                                                                       | 0/35 [00:00<?, ?it/s]
!!! Exception during processing !!!
Traceback (most recent call last):
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/execution.py", line 151, in recursive_execute
    output_data, output_ui = get_output_data(obj, input_data_all)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/execution.py", line 81, in get_output_data
    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/execution.py", line 74, in map_node_over_list
    results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/nodes.py", line 1206, in sample
    return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 170, in common_ksampler
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/AITemplate.py", line 304, in sample
    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)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/samplers.py", line 716, in sample
    samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/k_diffusion/sampling.py", line 580, in sample_dpmpp_2m
    denoised = model(x, sigmas[i] * s_in, **extra_args)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/samplers.py", line 319, in forward
    out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, cond_concat=cond_concat, model_options=model_options, seed=seed)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/k_diffusion/external.py", line 125, in forward
    eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/k_diffusion/external.py", line 151, in get_eps
    return self.inner_model.apply_model(*args, **kwargs)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/samplers.py", line 307, in apply_model
    out = sampling_function(self.inner_model.apply_model, x, timestep, uncond, cond, cond_scale, cond_concat, model_options=model_options, seed=seed)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/samplers.py", line 285, in sampling_function
    cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/comfy/samplers.py", line 262, in calc_cond_uncond_batch
    output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/inference.py", line 43, in apply_model
    return unet_inference(
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/inference.py", line 98, in unet_inference
    exe_module.run_with_tensors(inputs, ys, graph_mode=False)
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 565, in run_with_tensors
    outputs_ait = self.run(
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 468, in run
    return self._run_impl(
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 407, in _run_impl
    self.DLL.AITemplateModelContainerRun(
  File "/home/david/Development/machine-learning/stable-diffusion/ComfyUI/custom_nodes/AIT/AITemplate/ait/module/model.py", line 211, in _wrapped_func
    raise RuntimeError(f"Error in function: {method.__name__}")
RuntimeError: Error in function: AITemplateModelContainerRun