Zheng-Chong / CatVTON

CatVTON is a simple and efficient virtual try-on diffusion model with 1) Lightweight Network (899.06M parameters totally), 2) Parameter-Efficient Training (49.57M parameters trainable) and 3) Simplified Inference (< 8G VRAM for 1024X768 resolution).
Other
964 stars 114 forks source link

SCHP unable to run on CPU only environment #31

Closed kcltfask closed 3 months ago

kcltfask commented 3 months ago

Thanks for the great work. I am encountering a problem when running the script in a CPU-only environment (Colab with no GPU). Below are the error details:

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
[<ipython-input-22-144066cfdfe5>](https://localhost:8080/#) in <cell line: 5>()
      3 
      4 from utils import resize_and_crop
----> 5 from model.cloth_masker import AutoMasker as AM

9 frames
[/content/CatVTON/model/cloth_masker.py](https://localhost:8080/#) in <module>
      7 import torch
      8 
----> 9 from model.SCHP import SCHP  # type: ignore
     10 from model.DensePose import DensePose  # type: ignore
     11 

[/content/CatVTON/model/SCHP/__init__.py](https://localhost:8080/#) in <module>
----> 1 from model.SCHP import networks
      2 from model.SCHP.utils.transforms import get_affine_transform, transform_logits
      3 
      4 from collections import OrderedDict
      5 import torch

[/content/CatVTON/model/SCHP/networks/__init__.py](https://localhost:8080/#) in <module>
      1 from __future__ import absolute_import
      2 
----> 3 from model.SCHP.networks.AugmentCE2P import resnet101
      4 
      5 __factory = {

[/content/CatVTON/model/SCHP/networks/AugmentCE2P.py](https://localhost:8080/#) in <module>
     19 # Note here we adopt the InplaceABNSync implementation from https://github.com/mapillary/inplace_abn
     20 # By default, the InplaceABNSync module contains a BatchNorm Layer and a LeakyReLu layer
---> 21 from model.SCHP.modules import InPlaceABNSync
     22 
     23 BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')

[/content/CatVTON/model/SCHP/modules/__init__.py](https://localhost:8080/#) in <module>
----> 1 from .bn import ABN, InPlaceABN, InPlaceABNSync
      2 from .functions import ACT_RELU, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE
      3 from .misc import GlobalAvgPool2d, SingleGPU
      4 from .residual import IdentityResidualBlock
      5 from .dense import DenseModule

[/content/CatVTON/model/SCHP/modules/bn.py](https://localhost:8080/#) in <module>
      8     from Queue import Queue
      9 
---> 10 from .functions import *
     11 
     12 

[/content/CatVTON/model/SCHP/modules/functions.py](https://localhost:8080/#) in <module>
      8 
      9 _src_path = path.join(path.dirname(path.abspath(__file__)), "src")
---> 10 _backend = load(name="inplace_abn",
     11                 extra_cflags=["-O3"],
     12                 sources=[path.join(_src_path, f) for f in [

[/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py](https://localhost:8080/#) in load(name, sources, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, build_directory, verbose, with_cuda, is_python_module, is_standalone, keep_intermediates)
   1307         ...     verbose=True)
   1308     """
-> 1309     return _jit_compile(
   1310         name,
   1311         [sources] if isinstance(sources, str) else sources,

[/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py](https://localhost:8080/#) in _jit_compile(name, sources, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, build_directory, verbose, with_cuda, is_python_module, is_standalone, keep_intermediates)
   1743         return _get_exec_path(name, build_directory)
   1744 
-> 1745     return _import_module_from_library(name, build_directory, is_python_module)
   1746 
   1747 

[/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py](https://localhost:8080/#) in _import_module_from_library(module_name, path, is_python_module)
   2141         spec = importlib.util.spec_from_file_location(module_name, filepath)
   2142         assert spec is not None
-> 2143         module = importlib.util.module_from_spec(spec)
   2144         assert isinstance(spec.loader, importlib.abc.Loader)
   2145         spec.loader.exec_module(module)

ImportError: /tmp/inplace_abn/inplace_abn.so: cannot open shared object file: No such file or directory

I believe the issue arises because some dependencies of SCHP require CUDA and are only available in a CUDA environment. By the way, I have set export TORCH_EXTENSIONS_DIR=/tmp to overcome another issue, so you might see import errors from /tmp. Do you have a solution to run SCHP in a CPU-only environment?

Zheng-Chong commented 3 months ago

An edited version (has been set as default branch in this project) is just released, which modifies the SCHP code to avoid JIT and localize DensePose to avoid installation. It may be helpful for you to solve this issue, but I'm not sure if all modules can be run on CPU only.

kcltfask commented 3 months ago

Thank you for the update. After working around the changes, I am still encountering issues due to SCHP or DensePose when CUDA is not available. However, I have decided to give up on a CPU-only solution. Thanks again for your assistance.

Zheng-Chong commented 3 months ago

Using the CPU will significantly slow down the inference speed, generally consuming more than 10x of the time compared to a GPU. This means that a single try-on will take more than 5 minutes, which may be unacceptable for applications. However, you can use our Hugging Face Space or online demo.

kcltfask commented 3 months ago

Thank you for your advice, your online demo is fantastic.