app.py

import os


os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")
os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

try:
    import spaces
except ImportError:
    spaces = None

import gradio as gr


# Do not request a ZeroGPU worker during app startup. CUDA render extensions
# are installed from prebuilt wheels in requirements.txt.
APP_MODE = os.environ.get("DVD_SPACE_APP", "image").strip().lower()


def _maybe_preload(module, label: str):
    if os.environ.get("DVD_PRELOAD_DVD_GEN", "0").lower() in {"0", "false", "no"}:
        return
    module.log_event(f"preloading {label} DVD generation pipeline on cpu")
    module.ensure_dvd_gen_pipeline("cpu")
    module.log_event(f"{label} DVD generation pipeline preloaded on cpu")


if APP_MODE in {"image", "img"}:
    import app_dvd_image as app_module
    _maybe_preload(app_module, "image")
    demo = app_module.demo
elif APP_MODE in {"text", "txt"}:
    import app_dvd_text as app_module
    _maybe_preload(app_module, "text")
    demo = app_module.demo
elif APP_MODE in {"both", "all", "full"}:
    import app_dvd_image
    import app_dvd_text
    _maybe_preload(app_dvd_image, "image")
    _maybe_preload(app_dvd_text, "text")
    demo = gr.TabbedInterface(
        [app_dvd_image.demo, app_dvd_text.demo],
        ["Image", "Text"],
        title="DVD + TRELLIS Voxel Generation and Editing",
    )
else:
    raise ValueError("DVD_SPACE_APP must be 'image', 'text', or 'both'.")


if __name__ == "__main__":
    demo.queue().launch(show_api=False, show_error=True, ssr_mode=False)