Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """SplatAnything Gradio demo: python app.py (local GPU or CPU smoke run).""" | |
| import argparse | |
| import os | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parent | |
| def parse_args(argv=None): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument('--host', default=os.environ.get('HOST', '127.0.0.1')) | |
| parser.add_argument('--port', type=int, default=int(os.environ.get('PORT', '7860'))) | |
| parser.add_argument('--device', default=os.environ.get('SPLATANYTHING_DEVICE', 'cuda')) | |
| parser.add_argument('--config_path', default=os.environ.get('SPLATANYTHING_CONFIG', str(ROOT / 'configs/inference.yaml'))) | |
| parser.add_argument('--ckpt_path', default=os.environ.get('SPLATANYTHING_CKPT')) | |
| parser.add_argument('--geometry_ckpt_path', default=os.environ.get('WORLDMIRROR2_CKPT')) | |
| parser.add_argument('--sky_onnx_path', default=os.environ.get('SPLATANYTHING_SKY_CKPT')) | |
| parser.add_argument('--output_dir', default=os.environ.get('SPLATANYTHING_OUTPUT_DIR')) | |
| parser.add_argument('--examples_dir', default=os.environ.get('SPLATANYTHING_EXAMPLES_DIR', str(ROOT / 'examples'))) | |
| parser.add_argument('--install-viewer', action='store_true', help='Install the pinned npm viewer dependencies before launch.') | |
| args = parser.parse_args(argv) | |
| return args | |
| def build_app(args): | |
| # ZeroGPU must initialize CUDA emulation before Gradio or inference imports torch. | |
| zero = os.environ.get('SPLATANYTHING_ZEROGPU', os.environ.get('SPACES_ZERO_GPU', '0')).lower() in ('1', 'true') | |
| if zero: | |
| import spaces # noqa: F401 | |
| from demo.runtime import Pipeline | |
| from demo.ui import create_app | |
| pipeline = Pipeline(args.config_path, args.device, args.ckpt_path, args.geometry_ckpt_path, args.sky_onnx_path) | |
| predict_fn = None | |
| if zero: | |
| from demo.zero_gpu import initialize, predict | |
| initialize(pipeline) | |
| predict_fn = predict | |
| app = create_app(pipeline, args.output_dir, args.examples_dir, predict_fn=predict_fn) | |
| app.state.zero_gpu = zero | |
| return app | |
| def main(argv=None): | |
| args = parse_args(argv) | |
| if args.install_viewer: | |
| from demo.viewer import install_converter | |
| install_converter() | |
| app = build_app(args) | |
| import uvicorn | |
| if app.state.zero_gpu: | |
| # mount_gradio_app does not call Blocks.launch(), whose spaces SDK hook | |
| # packs preloaded CUDA tensors and reports the GPU functions to the host. | |
| # Expose the server as well so ZeroGPU workers close inherited listener FDs. | |
| server = uvicorn.Server(uvicorn.Config(app, host=args.host, port=args.port, access_log=False)) | |
| app.state.demo.server = server | |
| from demo.zero_gpu import startup | |
| startup() | |
| server.run() | |
| else: | |
| uvicorn.run(app, host=args.host, port=args.port, access_log=False) | |
| if __name__ == '__main__': | |
| main() | |