Update gradio_app/gradio_3dgen.py

gradio_app/gradio_3dgen.py

@@ -1,3 +1,5 @@
+# gradio_app/gradio_3dgen.py
+
 import spaces
 import os
 import gradio as gr
@@ -8,8 +10,12 @@ from gradio_app.custom_models.mvimg_prediction import run_mvprediction
 from gradio_app.custom_models.normal_prediction import predict_normals
 from scripts.refine_lr_to_sr import run_sr_fast
 from scripts.utils import save_glb_and_video
-# from scripts.multiview_inference import geo_reconstruct
 from scripts.multiview_inference import geo_reconstruct_part1, geo_reconstruct_part2, geo_reconstruct_part3
+import uuid 
+import shutil
+import logging 
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - UNIQUE3D - %(levelname)s - %(message)s')
 
 @spaces.GPU(duration=100)
 def run_mv(preview_img, input_processing, seed):
@@ -18,30 +24,57 @@ def run_mv(preview_img, input_processing, seed):
     rgb_pils, front_pil = run_mvprediction(preview_img, remove_bg=input_processing, seed=int(seed)) # 6s
     return rgb_pils, front_pil
 
-@spaces.GPU(duration=100) # seems split into multiple part will leads to `RuntimeError`, before fix it, still initialize here
+@spaces.GPU(duration=100)
 def generate3dv2(preview_img, input_processing, seed, render_video=True, do_refine=True, expansion_weight=0.1, init_type="std"):
-    if preview_img is None:
-        raise gr.Error("The input image is none!")
-    if isinstance(preview_img, str):
-        preview_img = Image.open(preview_img)
-    
-    rgb_pils, front_pil = run_mv(preview_img, input_processing, seed)
+    job_id = str(uuid.uuid4())
+    output_dir = os.path.join("/tmp/gradio/", job_id)
+    os.makedirs(output_dir, exist_ok=True)
+    logging.info(f"[{job_id}] Nueva generación iniciada. Directorio temporal creado en: {output_dir}")
 
-    vertices, faces, img_list = geo_reconstruct_part1(rgb_pils, None, front_pil, do_refine=do_refine, predict_normal=True, expansion_weight=expansion_weight, init_type=init_type)
+    try:
+        if preview_img is None:
+            raise gr.Error("The input image is none!")
+        if isinstance(preview_img, str):
+            preview_img = Image.open(preview_img)
+        
+        logging.info(f"[{job_id}] Ejecutando predicción multivista...")
+        rgb_pils, front_pil = run_mv(preview_img, input_processing, seed)
 
-    meshes = geo_reconstruct_part2(vertices, faces)
+        logging.info(f"[{job_id}] Reconstruyendo geometría (Parte 1)...")
+        vertices, faces, img_list = geo_reconstruct_part1(rgb_pils, None, front_pil, do_refine=do_refine, predict_normal=True, expansion_weight=expansion_weight, init_type=init_type)
 
-    new_meshes = geo_reconstruct_part3(meshes, img_list)
+        logging.info(f"[{job_id}] Reconstruyendo geometría (Parte 2)...")
+        meshes = geo_reconstruct_part2(vertices, faces)
 
-    vertices = new_meshes.verts_packed()
-    vertices = vertices / 2 * 1.35
-    vertices[..., [0, 2]] = - vertices[..., [0, 2]]
-    new_meshes = Meshes(verts=[vertices], faces=new_meshes.faces_list(), textures=new_meshes.textures)
-    
-    ret_mesh, video = save_glb_and_video("/tmp/gradio/generated", new_meshes, with_timestamp=True, dist=3.5, fov_in_degrees=2 / 1.35, cam_type="ortho", export_video=render_video)
-    return ret_mesh, video
+        logging.info(f"[{job_id}] Reconstruyendo geometría (Parte 3)...")
+        new_meshes = geo_reconstruct_part3(meshes, img_list)
+
+        vertices = new_meshes.verts_packed()
+        vertices = vertices / 2 * 1.35
+        vertices[..., [0, 2]] = - vertices[..., [0, 2]]
+        new_meshes = Meshes(verts=[vertices], faces=new_meshes.faces_list(), textures=new_meshes.textures)
+        
+        logging.info(f"[{job_id}] Guardando GLB y video en {output_dir}...")
+        ret_mesh, video = save_glb_and_video(output_dir, new_meshes, with_timestamp=False, dist=3.5, fov_in_degrees=2 / 1.35, cam_type="ortho", export_video=render_video, filename="model")
+        
+        logging.info(f"[{job_id}] Generación completada. Archivos resultantes: {ret_mesh}, {video}")
+        return ret_mesh, video
+
+    except Exception as e:
+        logging.error(f"[{job_id}] Ha ocurrido un error durante la generación: {e}", exc_info=True)
+        raise gr.Error(f"Error en la generación: {e}")
+        
+    finally:
+        logging.info(f"[{job_id}] Intentando limpiar el directorio temporal: {output_dir}")
+        if os.path.exists(output_dir):
+            try:
+                shutil.rmtree(output_dir)
+                logging.info(f"[{job_id}] Directorio temporal limpiado exitosamente.")
+            except Exception as e:
+                logging.error(f"[{job_id}] Error al limpiar el directorio temporal: {e}")
+        else:
+            logging.warning(f"[{job_id}] El directorio temporal no fue encontrado para la limpieza.")
 
-#######################################
 def create_ui(concurrency_id="wkl"):
     with gr.Row():
         with gr.Column(scale=1):
@@ -57,9 +90,7 @@ def create_ui(concurrency_id="wkl"):
                 examples_per_page=12
             )
             
-
         with gr.Column(scale=1):
-            # export mesh display
             output_mesh = gr.Model3D(value=None, label="Mesh Model", show_label=True, height=320, camera_position=(90, 90, 2))
             output_video = gr.Video(label="Preview", show_label=True, show_share_button=True, height=320, visible=False)
             
@@ -81,5 +112,5 @@ def create_ui(concurrency_id="wkl"):
         outputs=[output_mesh, output_video],
         concurrency_id=concurrency_id,
         api_name="generate3dv2",
-    ).success(clean_up, api_name=False)
+    )
     return input_image