Spaces:
Running on Zero
Running on Zero
| """Space entry point. Load both Hub models before serving GPU requests.""" | |
| import argparse | |
| import logging | |
| import os | |
| import spaces # Initialize ZeroGPU before importing PyTorch or Diffusers. | |
| from gradio_app import ( | |
| DEFAULT_MODEL, EXAMPLES, MODEL_IDS, UI_ASSETS, | |
| ModelRunner, build_app, load_pipeline, | |
| ) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--preview", action="store_true", help="Preview UI without downloading weights") | |
| parser.add_argument("--port", type=int, default=7860) | |
| args = parser.parse_args() | |
| logging.basicConfig(level=logging.INFO) | |
| model = os.environ.get("MOSS_MODEL_ID", DEFAULT_MODEL) | |
| if model not in MODEL_IDS: | |
| parser.error("MOSS_MODEL_ID must be OpenMOSS-Team/MOSS-Image-Base or OpenMOSS-Team/MOSS-Image-SFT") | |
| runner = ModelRunner("outputs/gradio") | |
| if not args.preview: | |
| shared_pipe = load_pipeline(model) | |
| runner.pipelines[model] = shared_pipe | |
| for repo in MODEL_IDS: | |
| if repo != model: | |
| runner.pipelines[repo] = load_pipeline(repo, shared_pipe=shared_pipe) | |
| if os.environ.get("MOSS_AOTI", "1") == "1": | |
| from aoti import prepare_aoti | |
| prepare_aoti(runner.pipelines.values()) | |
| demo = build_app(model=model, runner=runner) | |
| demo.launch( | |
| server_name="0.0.0.0", server_port=args.port, share=False, | |
| allowed_paths=[str(runner.output_dir), str(EXAMPLES.resolve()), str(UI_ASSETS.resolve())], | |
| ) | |
| if __name__ == "__main__": | |
| main() | |