integrate local trellis api into start_pipeline.py

2025-11-07 16:22:13 +02:00
parent 447449e1b3
commit 09f764c0df
6 changed files with 183 additions and 94 deletions
--- a/3d-generation-pipeline/generate_model_local.py
+++ b/3d-generation-pipeline/generate_model_local.py
@@ -0,0 +1,92 @@
+import subprocess
+import os
+import time
+import requests
+import base64
+
+from dotenv import load_dotenv
+
+load_dotenv()
+MODEL_FOLDER = os.environ["MODEL_FOLDER"]
+API_URL = os.environ["3D_GENERATION_URL"]
+
+
+def image_to_3d_subprocess(image_path, output_path):
+    venv_python = MODEL_FOLDER + r"\.venv\Scripts\python.exe"
+    script_path = MODEL_FOLDER + r"\run.py"
+
+    args = [image_path, "--output-dir", output_path]
+    command = [venv_python, script_path] + args
+
+    try:
+        # Run the subprocess
+        result = subprocess.run(command, capture_output=True, text=True)
+
+        # Print output and errors
+        print("STDOUT:\n", result.stdout)
+        print("STDERR:\n", result.stderr)
+        print("Return Code:", result.returncode)
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+
+
+def generate_no_preview(image_base64: str):
+    """Generate 3D model from a single base64-encoded image without previews.
+    
+    Args:
+        image_base64: Base64 string of the image (without 'data:image/...' prefix)
+    """
+    try:
+        # Set generation parameters
+        params = {
+            'image_base64': image_base64,
+            'seed': 42,
+            'ss_guidance_strength': 7.5,
+            'ss_sampling_steps': 30,
+            'slat_guidance_strength': 7.5,
+            'slat_sampling_steps': 30,
+            'mesh_simplify_ratio': 0.95,
+            'texture_size': 1024,
+            'output_format': 'glb'
+        }
+        
+        # Start generation
+        print("Starting generation...")
+        response = requests.post(f"{API_URL}/generate_no_preview", data=params)
+        response.raise_for_status()
+        
+        # Poll status until complete
+        while True:
+            status = requests.get(f"{API_URL}/status").json()
+            print(f"Progress: {status['progress']}%")
+            
+            if status['status'] == 'COMPLETE':
+                break
+            elif status['status'] == 'FAILED':
+                raise Exception(f"Generation failed: {status['message']}")
+            
+            time.sleep(1)
+        
+        # Download the model
+        print("Downloading model...")
+        response = requests.get(f"{API_URL}/download/model")
+        response.raise_for_status()
+        return response.content
+        
+        
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return None
+    
+def image_to_3d_api(image_path, output_path):
+    with open(image_path, 'rb') as image_file:
+        image_data = image_file.read()
+
+    base64_encoded = base64.b64encode(image_data).decode('utf-8')
+    model_binary = generate_no_preview(base64_encoded)
+
+    with open(output_path, 'wb') as f:
+        f.write(model_binary)
+