souging/TRELLIS_TextTo3D

1from typing import *2 3BACKEND = 'flash_attn' 4DEBUG = False5 6def __from_env():7    import os8    9    global BACKEND10    global DEBUG11    12    env_attn_backend = os.environ.get('ATTN_BACKEND')13    env_sttn_debug = os.environ.get('ATTN_DEBUG')14    15    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']:16        BACKEND = env_attn_backend17    if env_sttn_debug is not None:18        DEBUG = env_sttn_debug == '1'19 20    print(f"[ATTENTION] Using backend: {BACKEND}")21        22 23__from_env()24    25 26def set_backend(backend: Literal['xformers', 'flash_attn']):27    global BACKEND28    BACKEND = backend29 30def set_debug(debug: bool):31    global DEBUG32    DEBUG = debug33 34 35from .full_attn import *36from .modules import *37