Logging

app.py

@@ -17,6 +17,10 @@ from PIL import Image
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
+import logging
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - HF_SPACE_IMG - %(levelname)s - %(message)s')
+
 NUM_INFERENCE_STEPS = 8
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 # Constants
@@ -24,17 +28,29 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# Funciones auxiliares
 def start_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    logging.info(f"START SESSION: Creando directorio para la sesión {session_hash} en {user_dir}")
     os.makedirs(user_dir, exist_ok=True)
 
 def end_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    logging.info(f"END SESSION: Intentando eliminar el directorio de la sesión {session_hash} en {user_dir}")
+    if os.path.exists(user_dir):
+        try:
+            shutil.rmtree(user_dir)
+            logging.info(f"Directorio de la sesión {session_hash} eliminado correctamente.")
+        except Exception as e:
+            logging.error(f"Error al eliminar el directorio de la sesión {session_hash}: {e}")
+    else:
+        logging.warning(f"El directorio de la sesión {session_hash} no fue encontrado al intentar eliminarlo. Es posible que ya haya sido limpiado.")
 
 def preprocess_image(image: Image.Image) -> Image.Image:
+    logging.info("Preprocesando imagen para Trellis...")
     processed_image = trellis_pipeline.preprocess_image(image)
+    logging.info("Preprocesamiento de imagen completado.")
     return processed_image
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
@@ -74,7 +90,9 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
     return gs, mesh
 
 def get_seed(randomize_seed: bool, seed: int) -> int:
-    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
+    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
+    logging.info(f"Usando seed: {new_seed}")
+    return new_seed
 
 @spaces.GPU
 def generate_flux_image(
@@ -87,13 +105,15 @@ def generate_flux_image(
     req: gr.Request,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ) -> Image.Image:
-    """Generate image using Flux pipeline"""
+    session_hash = str(req.session_hash)
+    logging.info(f"[{session_hash}] Iniciando generate_flux_image con prompt: '{prompt[:50]}...'")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+        logging.info(f"[{session_hash}] Seed aleatorizado a: {seed}")
     generator = torch.Generator(device=device).manual_seed(seed)
-    prompt = "wbgmsst, " + prompt + ", 3D isometric, white background"
+    full_prompt = "wbgmsst, " + prompt + ", 3D isometric, white background"
     image = flux_pipeline(
-        prompt=prompt,
+        prompt=full_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=NUM_INFERENCE_STEPS,
         width=width,
@@ -101,13 +121,12 @@ def generate_flux_image(
         generator=generator,
     ).images[0]
     
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)  
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    unique_id = str(uuid.uuid4())[:8]
-    filename = f"{timestamp}_{unique_id}.png"
-    filepath = os.path.join(user_dir, filename)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    os.makedirs(user_dir, exist_ok=True)
+    
+    filepath = os.path.join(user_dir, "generated_2d_image.png")
     image.save(filepath)
+    logging.info(f"[{session_hash}] Imagen 2D guardada en: {filepath}")
     return image
 
 @spaces.GPU
@@ -120,7 +139,9 @@ def image_to_3d(
     slat_sampling_steps: int,
     req: gr.Request,
 ) -> Tuple[dict, str]:
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    session_hash = str(req.session_hash)
+    logging.info(f"[{session_hash}] Iniciando image_to_3d...")
+    user_dir = os.path.join(TMP_DIR, session_hash)
     outputs = trellis_pipeline.run(
         image,
         seed=seed,
@@ -135,6 +156,7 @@ def image_to_3d(
             "cfg_strength": slat_guidance_strength,
         },
     )
+    logging.info(f"[{session_hash}] Generación 3D completada. Renderizando video...")
     video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
     video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -142,6 +164,7 @@ def image_to_3d(
     imageio.mimsave(video_path, video, fps=15)
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
     torch.cuda.empty_cache()
+    logging.info(f"[{session_hash}] Video renderizado y estado empaquetado. Devolviendo: {video_path}")
     return state, video_path
 
 @spaces.GPU(duration=90)
@@ -151,12 +174,15 @@ def extract_glb(
     texture_size: int,
     req: gr.Request,
 ) -> Tuple[str, str]:
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    session_hash = str(req.session_hash)
+    logging.info(f"[{session_hash}] Iniciando extract_glb...")
+    user_dir = os.path.join(TMP_DIR, session_hash)
     gs, mesh = unpack_state(state)
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, 'sample.glb')
     glb.export(glb_path)
     torch.cuda.empty_cache()
+    logging.info(f"[{session_hash}] GLB extraído. Devolviendo: {glb_path}")
     return glb_path, glb_path
 
 # Interfaz Gradio
@@ -194,8 +220,7 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-    
-    # Variables adicionales para la generación 3D
+
     with gr.Accordion("3D Generation Settings", open=False):
         gr.Markdown("Stage 1: Sparse Structure Generation")
         with gr.Row():
@@ -205,15 +230,13 @@ with gr.Blocks() as demo:
         with gr.Row():
             slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
             slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-    
-    # Variables para la extracción de GLB
+
     with gr.Accordion("GLB Extraction Settings", open=False):
         mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
         texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
     
     output_buf = gr.State()
-    
-    # Event handlers
+
     demo.load(start_session)
     demo.unload(end_session)