Create app.py

app.py

@@ -0,0 +1,98 @@
+import gradio as gr
+import sqlite3
+import threading
+import time
+import requests
+import os
+import shutil
+
+# Use persistent storage
+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()
+
+# 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'/data/outputs/result_{job_id}.glb'
+    os.makedirs('/data/outputs', exist_ok=True)
+    with open(output_path, 'wb') as f:
+        f.write(response.content)
+    return output_path
+
+# Processing logic (same as before)
+def process_job(job_id):
+    try:
+        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))
+        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()
+
+# Gradio interface (same as before)
+def submit_images(files):
+    if not files:
+        return "No files uploaded."
+    for file in files:
+        c.execute("INSERT INTO jobs (status) VALUES ('submitted')")
+        job_id = c.lastrowid
+        conn.commit()
+        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()
+        threading.Thread(target=process_job, args=(job_id,), daemon=True).start()
+    return "Jobs submitted. Check the status tab."
+
+def get_status():
+    c.execute("SELECT id, image_path, status, output_path FROM jobs")
+    return c.fetchall()
+
+with gr.Blocks(title="TRELLIS 3D Generator") as demo:
+    with gr.Tab("Upload"):
+        files_input = gr.File(file_count="multiple", label="Upload Images (JPG/PNG)")
+        submit_btn = gr.Button("Submit")
+        output_msg = gr.Textbox(label="Message")
+        submit_btn.click(fn=submit_images, inputs=files_input, outputs=output_msg)
+    with gr.Tab("Status"):
+        status_table = gr.DataFrame(
+            headers=["ID", "Image Path", "Status", "Output Path"],
+            label="Job Status"
+        )
+        refresh_btn = gr.Button("Refresh")
+        refresh_btn.click(fn=get_status, inputs=None, outputs=status_table)
+
+demo.launch()