Update app.py

app.py

@@ -2,69 +2,81 @@ import gradio as gr
 import sqlite3
 import threading
 import time
-import requests
 import os
 import shutil
-import spaces
+from gradio_client import Client, handle_file
 
+# Ensure /data/ directory exists
+os.makedirs('/data', exist_ok=True)
 
-# Use persistent storage
-conn = sqlite3.connect('/tmp/jobs.db', check_same_thread=False)
+# Database setup
+conn = sqlite3.connect('/data/jobs.db', check_same_thread=False)
 c = conn.cursor()
 c.execute('''CREATE TABLE IF NOT EXISTS jobs
              (id INTEGER PRIMARY KEY, image_path TEXT, job_id TEXT, status TEXT, output_path TEXT)''')
 conn.commit()
 
-@spaces.GPU
-# API functions (adjust endpoints per TRELLIS API docs)
-def upload_image(image_path):
-    with open(image_path, 'rb') as f:
-        response = requests.post(
-            'https://huggingface.co/spaces/jkorstad/TRELLIS/api/upload',
-            files={'file': f}
-        )
-        response.raise_for_status()
-    return response.json()['job_id']
-
-def check_status(job_id):
-    response = requests.get(f'https://huggingface.co/spaces/jkorstad/TRELLIS/api/status/{job_id}')
-    response.raise_for_status()
-    return response.json()['status']
-
-def get_result(job_id):
-    response = requests.get(f'https://huggingface.co/spaces/jkorstad/TRELLIS/api/result/{job_id}')
-    response.raise_for_status()
-    output_path = f'/tmp/outputs/result_{job_id}.glb'
-    os.makedirs('/tmp/outputs', exist_ok=True)
-    with open(output_path, 'wb') as f:
-        f.write(response.content)
-    return output_path
+# TRELLIS API client
+trellis_client = Client("jkorstad/TRELLIS")
 
-# Processing logic (same as before)
+# Processing logic with three-step TRELLIS workflow
 def process_job(job_id):
     try:
+        # Get the uploaded image path
         c.execute("SELECT image_path FROM jobs WHERE id=?", (job_id,))
         image_path = c.fetchone()[0]
-        api_job_id = upload_image(image_path)
-        c.execute("UPDATE jobs SET job_id=?, status='processing' WHERE id=?", (api_job_id, job_id))
+        
+        # Step 1: Preprocess the image
+        c.execute("UPDATE jobs SET status='preprocessing' WHERE id=?", (job_id,))
+        conn.commit()
+        preprocessed_image = trellis_client.predict(
+            image=handle_file(image_path),
+            api_name="/preprocess_image"
+        )
+        
+        # Step 2: Generate 3D asset
+        c.execute("UPDATE jobs SET status='generating' WHERE id=?", (job_id,))
+        conn.commit()
+        time.sleep(5)  # Wait between steps; adjust based on observed timing
+        result_3d = trellis_client.predict(
+            image=handle_file(preprocessed_image),  # Use preprocessed image
+            multiimages=[],
+            seed=0,  # Default; could make configurable
+            ss_guidance_strength=7.5,
+            ss_sampling_steps=12,
+            slat_guidance_strength=3,
+            slat_sampling_steps=12,
+            multiimage_algo="stochastic",
+            api_name="/image_to_3d"
+        )
+        video_path = result_3d['video']  # Extract video filepath from dict
+        
+        # Step 3: Extract GLB
+        c.execute("UPDATE jobs SET status='extracting' WHERE id=?", (job_id,))
+        conn.commit()
+        time.sleep(10)  # Wait for 3D processing; adjust as needed
+        glb_result = trellis_client.predict(
+            mesh_simplify=0.95,
+            texture_size=1024,
+            api_name="/extract_glb"
+        )
+        glb_path = glb_result[0]  # First element is the GLB filepath
+        
+        # Move GLB to persistent storage
+        output_path = f'/data/outputs/result_{job_id}.glb'
+        os.makedirs('/data/outputs', exist_ok=True)
+        shutil.move(glb_path, output_path)
+        
+        # Update job status
+        c.execute("UPDATE jobs SET status='completed', output_path=? WHERE id=?", (output_path, job_id))
         conn.commit()
-        while True:
-            status = check_status(api_job_id)
-            if status == 'completed':
-                output_path = get_result(api_job_id)
-                c.execute("UPDATE jobs SET status='completed', output_path=? WHERE id=?", (output_path, job_id))
-                conn.commit()
-                break
-            elif status == 'failed':
-                c.execute("UPDATE jobs SET status='failed' WHERE id=?", (job_id,))
-                conn.commit()
-                break
-            time.sleep(5)
+        
     except Exception as e:
         c.execute("UPDATE jobs SET status='failed' WHERE id=?", (job_id,))
         conn.commit()
+        print(f"Error processing job {job_id}: {e}")
 
-# Gradio interface (same as before)
+# Gradio interface
 def submit_images(files):
     if not files:
         return "No files uploaded."
@@ -72,8 +84,8 @@ def submit_images(files):
         c.execute("INSERT INTO jobs (status) VALUES ('submitted')")
         job_id = c.lastrowid
         conn.commit()
-        image_path = f'/tmp/inputs/input_{job_id}.jpg'
-        os.makedirs('/tmp/inputs', exist_ok=True)
+        image_path = f'/data/inputs/input_{job_id}.jpg'
+        os.makedirs('/data/inputs', exist_ok=True)
         shutil.copy(file.name, image_path)
         c.execute("UPDATE jobs SET image_path=? WHERE id=?", (image_path, job_id))
         conn.commit()