Spaces:
Running
Running
File size: 2,663 Bytes
f3b96de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
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] |