Spaces:
Running
Running
| import subprocess | |
| import os | |
| from pathlib import Path | |
| import glob | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| def run_inference(image_path: str, output_dir: str): | |
| """ | |
| Run Stable Fast 3D inference on the input image | |
| Args: | |
| image_path: Path to the input image | |
| output_dir: Directory to save output files | |
| Returns: | |
| Path to the generated GLB file | |
| """ | |
| # Get the repository path | |
| repo_path = Path("/app/stable-fast-3d") | |
| logger.info(f"Running inference on image: {image_path}") | |
| logger.info(f"Output directory: {output_dir}") | |
| # Run the model | |
| command = [ | |
| "python", "run.py", | |
| image_path, | |
| "--output-dir", output_dir, | |
| "--pretrained-model", "stabilityai/stable-fast-3d", | |
| "--texture-resolution", "2048", | |
| "--remesh_option", "quad" | |
| ] | |
| logger.info(f"Running command: {' '.join(command)}") | |
| print(f"Reached at subprocess.run() The output path is {output_dir}") | |
| try: | |
| HF_CACHE_PATH = "/tmp/huggingface_cache" | |
| os.makedirs(HF_CACHE_PATH, exist_ok=True) | |
| env = os.environ.copy() | |
| env["HF_HOME"] = "/tmp/huggingface" | |
| env["TRANSFORMERS_CACHE"] = HF_CACHE_PATH | |
| env["HF_HUB_CACHE"] = HF_CACHE_PATH | |
| env["NUMBA_CACHE_DIR"] = "/tmp" | |
| result = subprocess.run( | |
| command, | |
| cwd=repo_path, | |
| check=True, | |
| capture_output=True, | |
| text=True, | |
| env=env) # pass updated env here | |
| logger.info(f"Subprocess STDOUT:\n{result.stdout}") | |
| logger.info(f"Subprocess STDERR:\n{result.stderr}") | |
| except subprocess.CalledProcessError as e: | |
| logger.error(f"Subprocess failed with exit code {e.returncode}") | |
| logger.error(f"STDOUT:\n{e.stdout}") | |
| logger.error(f"STDERR:\n{e.stderr}") | |
| raise Exception(f"Subprocess failed: {e.stderr}") | |
| # Find the generated GLB file | |
| # The model typically outputs files to a subdirectory named after the input image | |
| # without the file extension | |
| base_name = os.path.basename(image_path) | |
| file_name_without_ext = os.path.splitext(base_name)[0] | |
| # Look for GLB files in the output directory and its subdirectories | |
| glb_files = glob.glob(os.path.join(output_dir, "**", "*.glb"), recursive=True) | |
| if not glb_files: | |
| logger.error(f"No GLB files found in {output_dir}") | |
| raise Exception("No GLB file was generated") | |
| logger.info(f"Found GLB file: {glb_files[0]}") | |
| print(f"Returned the path {glb_files[0]}") | |
| return glb_files[0] |