Update app.py

app.py

@@ -52,9 +52,10 @@ HEADER = """
 ## 🚀 Launch Your Creation:
 1. **Upload an Image** (clear, single-object images shine brightest)
 2. **Choose a Style Filter** to infuse your unique vision
-3. Click **Generate 3D Model** to sculpt your mesh
-4. Click **Apply Texture** to bring your model to life
-5. **Download GLB** to share your masterpiece
+3. **Select Output Format** (OBJ or GLB)
+4. Click **Generate 3D Model** to sculpt your mesh
+5. Click **Apply Texture** to bring your model to life
+6. **Download** your creation in your chosen format
 <p style="font-size: 0.9em; margin-top: 10px; color: #D1D5DB;">Powered by cutting-edge AI and multi-view technology from AnvilInteractive Solutions. Join our <a href="https://www.anvilinteractive.com/community" style="color: #A78BFA; text-decoration: none;">PolyGenixAI Community</a> to connect with creators and spark inspiration.</p>
 <style>
   @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
@@ -192,7 +193,7 @@ def get_random_seed(randomize_seed, seed):
     return seed
 
 @spaces.GPU(duration=180)
-def run_full(image: str, req: gr.Request):
+def run_full(image: str, output_format: str, req: gr.Request):
     seed = 0
     num_inference_steps = 50
     guidance_scale = 7.5
@@ -217,8 +218,9 @@ def run_full(image: str, req: gr.Request):
     
     save_dir = os.path.join(TMP_DIR, "examples")
     os.makedirs(save_dir, exist_ok=True)
-    mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
-    mesh.export(mesh_path)
+    file_extension = output_format.lower()
+    mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.{file_extension}")
+    mesh.export(mesh_path, file_type=file_extension)
     print("save to ", mesh_path)
 
     torch.cuda.empty_cache()
@@ -295,10 +297,10 @@ def run_full(image: str, req: gr.Request):
         device=DEVICE,
     )
 
-    textured_glb_path = texture_pipe(
+    textured_mesh_path = texture_pipe(
         mesh_path=mesh_path,
         save_dir=save_dir,
-        save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
+        save_name=f"polygenixai_texture_mesh_{get_random_hex()}.{file_extension}",
         uv_unwarp=True,
         uv_size=4096,
         rgb_path=mv_image_path,
@@ -306,8 +308,7 @@ def run_full(image: str, req: gr.Request):
         camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
     )
 
-    return image_seg, mesh_path, textured_glb_path
-    
+    return image_seg, mesh_path, textured_mesh_path
 
 @spaces.GPU()
 @torch.no_grad()
@@ -324,6 +325,7 @@ def image_to_3d(
     guidance_scale: float,
     simplify: bool,
     target_face_num: int,
+    output_format: str,
     req: gr.Request
 ):
     outputs = triposg_pipe(
@@ -341,8 +343,9 @@ def image_to_3d(
         mesh = simplify_mesh(mesh, target_face_num)
     
     save_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
-    mesh.export(mesh_path)
+    file_extension = output_format.lower()
+    mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.{file_extension}")
+    mesh.export(mesh_path, file_type=file_extension)
     print("save to ", mesh_path)
 
     torch.cuda.empty_cache()
@@ -351,7 +354,7 @@ def image_to_3d(
 
 @spaces.GPU(duration=120)
 @torch.no_grad()
-def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
+def run_texture(image: Image, mesh_path: str, seed: int, output_format: str, req: gr.Request):
     height, width = 768, 768
     # Prepare cameras
     cameras = get_orthogonal_camera(
@@ -425,10 +428,10 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
         device=DEVICE,
     )
 
-    textured_glb_path = texture_pipe(
+    textured_mesh_path = texture_pipe(
         mesh_path=mesh_path,
         save_dir=save_dir,
-        save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
+        save_name=f"polygenixai_texture_mesh_{get_random_hex()}.{output_format.lower()}",
         uv_unwarp=True,
         uv_size=4096,
         rgb_path=mv_image_path,
@@ -436,7 +439,7 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
         camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
     )
 
-    return textured_glb_path
+    return textured_mesh_path
 
 
 with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr-panel { background-color: #2D2D2D; }") as demo:
@@ -456,6 +459,13 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
                             info="Select a style to inspire your 3D model (optional)",
                             elem_classes="gr-dropdown"
                         )
+                        output_format = gr.Dropdown(
+                            choices=["GLB", "OBJ"],
+                            label="Output Format",
+                            value="GLB",
+                            info="Choose the format for your 3D model",
+                            elem_classes="gr-dropdown"
+                        )
                         seed = gr.Slider(
                             label="Seed",
                             minimum=0,
@@ -498,7 +508,7 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
                 with gr.Column(scale=1):
                     model_output = gr.Model3D(label="3D Model Preview", interactive=False, height=400, elem_classes="gr-panel")
                     textured_model_output = gr.Model3D(label="Textured 3D Model", interactive=False, height=400, elem_classes="gr-panel")
-                    download_button = gr.Button("Download GLB", variant="secondary", elem_classes="gr-button-secondary")
+                    download_button = gr.Button("Download Model", variant="secondary", elem_classes="gr-button-secondary")
 
         with gr.Tab("Cosmic Gallery"):
             gr.Markdown("### Discover Stellar Creations")
@@ -508,7 +518,7 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
                     for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
                 ],
                 fn=run_full,
-                inputs=[image_prompts],
+                inputs=[image_prompts, output_format],
                 outputs=[seg_image, model_output, textured_model_output],
                 cache_examples=True,
             )
@@ -530,14 +540,15 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
             num_inference_steps,
             guidance_scale,
             reduce_face,
-            target_face_num
+            target_face_num,
+            output_format
         ],
         outputs=[model_output]
     ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
 
     gen_texture_button.click(
         run_texture,
-        inputs=[image_prompts, model_output, seed],
+        inputs=[image_prompts, model_output, seed, output_format],
         outputs=[textured_model_output]
     )