Spaces:

---

Duplicated from  microsoft/TRELLIS

broyang

/

3dai

Like 36

Running on Zero

from functools import partial

import contextvars

import threading

import html

import re

import argparse

import base64

import inspect

import importlib.util

import json

import os

import sys

import time

from datetime import datetime

from pathlib import Path

from uuid import uuid4

from typing import Any, Dict, Generator, List, Optional, Tuple

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

RUNNING_ON_SPACES = bool(os.getenv("SPACE_ID"))

if RUNNING_ON_SPACES and importlib.util.find_spec("xformers") is not None:

_attn_backend = "xformers"

elif os.getenv("SPACE_ID"):

_attn_backend = "sdpa"

elif importlib.util.find_spec("flash_attn") is not None:

_attn_backend = "flash_attn"

elif importlib.util.find_spec("xformers") is not None:

_attn_backend = "xformers"

else:

_attn_backend = "sdpa"

os.environ.setdefault("ATTN_BACKEND", _attn_backend)

os.environ.setdefault("SPARSE_ATTN_BACKEND", _attn_backend)

os.environ.setdefault("SPARSE_ATTN", _attn_backend)

os.environ.setdefault("SPCONV_ALGO", "native")

os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(

os.path.dirname(os.path.abspath(__file__)), "autotune_cache.json"

)

os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = "1"

os.environ.setdefault("TRELLIS_REMBG_MODEL", "briaai/RMBG-2.0")

import cv2

import gradio as gr

import imageio

import imageio_ffmpeg

import numpy as np

import spaces

import logging

from bake_flight import BakeFlights

from quota_fallback import (

ALLOWANCE_ERROR,

CAPACITY_ERROR,

EXPORT,

GENERATE,

WAITING_STATUS,

FallbackCoordinator,

PendingCache,

is_eligible_embed,

is_quota_denial,

request_header,

safe_message,

)

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

logger = logging.getLogger(__name__)

FALLBACK = FallbackCoordinator(PendingCache(), {})

FLIGHTS = BakeFlights()

def configure_xformers_attention() -> None:

global _attn_backend

if not RUNNING_ON_SPACES or _attn_backend != "xformers":

return

try:

import xformers.ops as xformers_ops

original_memory_efficient_attention = xformers_ops.memory_efficient_attention

cutlass_attention_ops = (

xformers_ops.fmha.cutlass.FwOp,

xformers_ops.fmha.cutlass.BwOp,

)

except (ImportError, AttributeError, OSError) as error:

logger.warning("xformers unavailable; falling back to SDPA: %s", error)

_attn_backend = "sdpa"

os.environ["ATTN_BACKEND"] = _attn_backend

os.environ["SPARSE_ATTN_BACKEND"] = _attn_backend

os.environ["SPARSE_ATTN"] = _attn_backend

return

def memory_efficient_attention_with_cutlass(*args, **kwargs):

kwargs.setdefault("op", cutlass_attention_ops)

return original_memory_efficient_attention(*args, **kwargs)

xformers_ops.memory_efficient_attention = memory_efficient_attention_with_cutlass

configure_xformers_attention()