jkorstad commited on
Commit
24a4bd7
·
verified ·
1 Parent(s): bb01f49

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
app.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import sqlite3
3
+ import threading
4
+ import time
5
+ import requests
6
+ import os
7
+ import shutil
8
+
9
+ # Use persistent storage
10
+ conn = sqlite3.connect('/data/jobs.db', check_same_thread=False)
11
+ c = conn.cursor()
12
+ c.execute('''CREATE TABLE IF NOT EXISTS jobs
13
+ (id INTEGER PRIMARY KEY, image_path TEXT, job_id TEXT, status TEXT, output_path TEXT)''')
14
+ conn.commit()
15
+
16
+ # API functions (adjust endpoints per TRELLIS API docs)
17
+ def upload_image(image_path):
18
+ with open(image_path, 'rb') as f:
19
+ response = requests.post(
20
+ 'https://huggingface.co/spaces/jkorstad/TRELLIS/api/upload',
21
+ files={'file': f}
22
+ )
23
+ response.raise_for_status()
24
+ return response.json()['job_id']
25
+
26
+ def check_status(job_id):
27
+ response = requests.get(f'https://huggingface.co/spaces/jkorstad/TRELLIS/api/status/{job_id}')
28
+ response.raise_for_status()
29
+ return response.json()['status']
30
+
31
+ def get_result(job_id):
32
+ response = requests.get(f'https://huggingface.co/spaces/jkorstad/TRELLIS/api/result/{job_id}')
33
+ response.raise_for_status()
34
+ output_path = f'/data/outputs/result_{job_id}.glb'
35
+ os.makedirs('/data/outputs', exist_ok=True)
36
+ with open(output_path, 'wb') as f:
37
+ f.write(response.content)
38
+ return output_path
39
+
40
+ # Processing logic (same as before)
41
+ def process_job(job_id):
42
+ try:
43
+ c.execute("SELECT image_path FROM jobs WHERE id=?", (job_id,))
44
+ image_path = c.fetchone()[0]
45
+ api_job_id = upload_image(image_path)
46
+ c.execute("UPDATE jobs SET job_id=?, status='processing' WHERE id=?", (api_job_id, job_id))
47
+ conn.commit()
48
+ while True:
49
+ status = check_status(api_job_id)
50
+ if status == 'completed':
51
+ output_path = get_result(api_job_id)
52
+ c.execute("UPDATE jobs SET status='completed', output_path=? WHERE id=?", (output_path, job_id))
53
+ conn.commit()
54
+ break
55
+ elif status == 'failed':
56
+ c.execute("UPDATE jobs SET status='failed' WHERE id=?", (job_id,))
57
+ conn.commit()
58
+ break
59
+ time.sleep(5)
60
+ except Exception as e:
61
+ c.execute("UPDATE jobs SET status='failed' WHERE id=?", (job_id,))
62
+ conn.commit()
63
+
64
+ # Gradio interface (same as before)
65
+ def submit_images(files):
66
+ if not files:
67
+ return "No files uploaded."
68
+ for file in files:
69
+ c.execute("INSERT INTO jobs (status) VALUES ('submitted')")
70
+ job_id = c.lastrowid
71
+ conn.commit()
72
+ image_path = f'/data/inputs/input_{job_id}.jpg'
73
+ os.makedirs('/data/inputs', exist_ok=True)
74
+ shutil.copy(file.name, image_path)
75
+ c.execute("UPDATE jobs SET image_path=? WHERE id=?", (image_path, job_id))
76
+ conn.commit()
77
+ threading.Thread(target=process_job, args=(job_id,), daemon=True).start()
78
+ return "Jobs submitted. Check the status tab."
79
+
80
+ def get_status():
81
+ c.execute("SELECT id, image_path, status, output_path FROM jobs")
82
+ return c.fetchall()
83
+
84
+ with gr.Blocks(title="TRELLIS 3D Generator") as demo:
85
+ with gr.Tab("Upload"):
86
+ files_input = gr.File(file_count="multiple", label="Upload Images (JPG/PNG)")
87
+ submit_btn = gr.Button("Submit")
88
+ output_msg = gr.Textbox(label="Message")
89
+ submit_btn.click(fn=submit_images, inputs=files_input, outputs=output_msg)
90
+ with gr.Tab("Status"):
91
+ status_table = gr.DataFrame(
92
+ headers=["ID", "Image Path", "Status", "Output Path"],
93
+ label="Job Status"
94
+ )
95
+ refresh_btn = gr.Button("Refresh")
96
+ refresh_btn.click(fn=get_status, inputs=None, outputs=status_table)
97
+
98
+ demo.launch()