Update app.py

app.py

@@ -7,13 +7,67 @@ import os
 import io
 import json
 import tempfile
+import logging
+from datetime import datetime
+import hashlib
+import shutil
 import gradio as gr
 from google import genai
 from google.genai import types
 from PIL import Image
 import modal
 
+# ============================================================
+# LOGGING SETUP
+# ============================================================
+
+# Setup logging for console output (visible in HF Spaces logs)
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler()  # This goes to HF Spaces logs
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Create persistent logs directory (use HF Datasets for permanent storage)
+LOGS_DIR = "generation_logs"
+os.makedirs(LOGS_DIR, exist_ok=True)
+
+def save_generation_log(prompt, images_dict, metadata=None):
+    """Save generation logs with images"""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    session_id = hashlib.md5(f"{prompt}{timestamp}".encode()).hexdigest()[:8]
+    
+    session_dir = os.path.join(LOGS_DIR, f"{timestamp}_{session_id}")
+    os.makedirs(session_dir, exist_ok=True)
+    
+    # Save metadata
+    log_data = {
+        "timestamp": timestamp,
+        "session_id": session_id,
+        "prompt": prompt,
+        "metadata": metadata or {}
+    }
+    
+    log_file = os.path.join(session_dir, "metadata.json")
+    with open(log_file, 'w') as f:
+        json.dump(log_data, f, indent=2)
+    
+    # Save images
+    for name, image_path in images_dict.items():
+        if image_path and os.path.exists(image_path):
+            dest = os.path.join(session_dir, f"{name}.png")
+            shutil.copy(image_path, dest)
+    
+    logger.info(f"✓ Saved logs to: {session_dir}")
+    return session_dir
+
+# ============================================================
 # Initialize Gemini client
+# ============================================================
+
 client = None
 
 def init_gemini():
@@ -22,7 +76,9 @@ def init_gemini():
     if api_key:
         os.environ["GEMINI_API_KEY"] = api_key
         client = genai.Client()
+        logger.info("✓ Gemini client initialized")
         return True
+    logger.error("✗ GEMINI_API_KEY not found")
     return False
 
 def image_to_bytes(image):
@@ -33,6 +89,7 @@ def image_to_bytes(image):
 
 def run_sam3d(image, mask):
     """Send image and mask to SAM-3D on Modal"""
+    logger.info("→ Sending to SAM-3D on Modal...")
     img_bytes = image_to_bytes(image.convert("RGB"))
     mask_bytes = image_to_bytes(mask)
     
@@ -40,6 +97,7 @@ def run_sam3d(image, mask):
     model = SAM3DModel()
     ply_bytes, glb_bytes = model.reconstruct.remote(img_bytes, mask_bytes)
     
+    logger.info(f"✓ SAM-3D complete - PLY: {len(ply_bytes)} bytes, GLB: {len(glb_bytes) if glb_bytes else 0} bytes")
     return ply_bytes, glb_bytes
 
 
@@ -57,6 +115,9 @@ def generate_3d_model(prompt: str) -> str:
     Returns:
         JSON string with paths to generated files
     """
+    logger.info(f"=== MCP TOOL: generate_3d_model ===")
+    logger.info(f"Prompt: {prompt}")
+    
     if not client:
         if not init_gemini():
             return json.dumps({"error": "GEMINI_API_KEY not configured"})
@@ -64,6 +125,7 @@ def generate_3d_model(prompt: str) -> str:
     try:
         # STEP 1: Generate image
         initial_prompt = f"{prompt}, three-quarter front view angle, natural daylight, soft shadows showing depth and contours, clean simple background, full object visible, photorealistic"
+        logger.info("→ Generating initial image...")
         
         response_gen = client.models.generate_content(
             model="gemini-2.5-flash-image",
@@ -78,9 +140,13 @@ def generate_3d_model(prompt: str) -> str:
                 break
         
         if initial_image is None:
+            logger.error("✗ Image generation failed")
             return json.dumps({"error": "Image generation failed"})
         
+        logger.info(f"✓ Initial image generated: {initial_image.size}")
+        
         # STEP 2: Remove background
+        logger.info("→ Removing background...")
         edit_prompt = "Remove the background completely, make the background transparent. Preserve the object's shadow for realism."
         image_part = types.Part.from_bytes(
             data=image_to_bytes(initial_image), 
@@ -100,8 +166,11 @@ def generate_3d_model(prompt: str) -> str:
                 break
         
         if final_image is None:
+            logger.error("✗ Background removal failed")
             return json.dumps({"error": "Background removal failed"})
         
+        logger.info("✓ Background removed")
+        
         # STEP 3: Create grayscale mask
         gray = final_image.convert("L")
         
@@ -129,6 +198,22 @@ def generate_3d_model(prompt: str) -> str:
             with open(glb_path, 'wb') as f:
                 f.write(glb_bytes)
         
+        # Save logs
+        images_dict = {
+            "original": original_path,
+            "transparent": nobg_path,
+            "mask": mask_path
+        }
+        metadata = {
+            "type": "generation",
+            "has_glb": glb_path is not None,
+            "ply_size_bytes": len(ply_bytes),
+            "glb_size_bytes": len(glb_bytes) if glb_bytes else 0
+        }
+        log_dir = save_generation_log(prompt, images_dict, metadata)
+        
+        logger.info(f"✓ Generation complete!")
+        
         return json.dumps({
             "success": True,
             "prompt": prompt,
@@ -137,10 +222,12 @@ def generate_3d_model(prompt: str) -> str:
             "mask_image": mask_path,
             "ply_model": ply_path,
             "glb_model": glb_path,
+            "log_directory": log_dir,
             "message": f"Successfully generated 3D model for: {prompt}"
         })
         
     except Exception as e:
+        logger.error(f"✗ Error: {e}", exc_info=True)
         return json.dumps({"error": str(e)})
 
 
@@ -155,12 +242,16 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
     Returns:
         JSON string with paths to the new edited files
     """
+    logger.info(f"=== MCP TOOL: edit_3d_model ===")
+    logger.info(f"Edit: {edit_prompt}")
+    
     if not client:
         if not init_gemini():
             return json.dumps({"error": "GEMINI_API_KEY not configured"})
     
     try:
         current_image = Image.open(transparent_image_path)
+        logger.info(f"→ Loaded image: {current_image.size}")
         
         image_part = types.Part.from_bytes(
             data=image_to_bytes(current_image), 
@@ -168,6 +259,7 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
         )
         
         full_edit_prompt = f"{edit_prompt}. Keep the background transparent. Maintain image quality and lighting."
+        logger.info("→ Applying edit...")
         
         response_edit = client.models.generate_content(
             model="gemini-3-pro-image-preview",
@@ -182,8 +274,11 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
                 break
         
         if edited_image is None:
+            logger.error("✗ Edit failed")
             return json.dumps({"error": "Edit failed"})
         
+        logger.info("✓ Edit applied")
+        
         gray = edited_image.convert("L")
         ply_bytes, glb_bytes = run_sam3d(edited_image, gray)
         
@@ -205,6 +300,19 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
             with open(glb_path, 'wb') as f:
                 f.write(glb_bytes)
         
+        # Save logs
+        images_dict = {
+            "edited": nobg_path,
+            "mask": mask_path
+        }
+        metadata = {
+            "type": "edit",
+            "has_glb": glb_path is not None
+        }
+        log_dir = save_generation_log(edit_prompt, images_dict, metadata)
+        
+        logger.info(f"✓ Edit complete!")
+        
         return json.dumps({
             "success": True,
             "edit_prompt": edit_prompt,
@@ -212,10 +320,12 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
             "mask_image": mask_path,
             "ply_model": ply_path,
             "glb_model": glb_path,
+            "log_directory": log_dir,
             "message": f"Successfully applied edit: {edit_prompt}"
         })
         
     except Exception as e:
+        logger.error(f"✗ Error: {e}", exc_info=True)
         return json.dumps({"error": str(e)})
 
 
@@ -225,6 +335,9 @@ def edit_3d_model(edit_prompt: str, transparent_image_path: str) -> str:
 
 def generate_3d_ui(prompt, progress=gr.Progress()):
     """UI wrapper with progress updates"""
+    logger.info(f"=== NEW GENERATION REQUEST ===")
+    logger.info(f"Prompt: {prompt}")
+    
     if not client:
         if not init_gemini():
             raise gr.Error("GEMINI_API_KEY not set in Space secrets")
@@ -248,8 +361,11 @@ def generate_3d_ui(prompt, progress=gr.Progress()):
         
         if initial_image is None:
             raise gr.Error("Image generation failed")
+        
+        logger.info(f"✓ Image generated: {initial_image.size}")
             
     except Exception as e:
+        logger.error(f"✗ Image generation failed: {e}")
         raise gr.Error(f"Image generation failed: {e}")
     
     progress(0.3, desc="Removing background...")
@@ -274,8 +390,11 @@ def generate_3d_ui(prompt, progress=gr.Progress()):
         
         if final_image is None:
             raise gr.Error("Background removal failed")
+        
+        logger.info("✓ Background removed")
             
     except Exception as e:
+        logger.error(f"✗ Background removal failed: {e}")
         raise gr.Error(f"Background removal failed: {e}")
     
     progress(0.4, desc="Creating mask...")
@@ -286,6 +405,7 @@ def generate_3d_ui(prompt, progress=gr.Progress()):
     try:
         ply_bytes, glb_bytes = run_sam3d(final_image, gray)
     except Exception as e:
+        logger.error(f"✗ SAM-3D failed: {e}")
         raise gr.Error(f"SAM-3D failed: {e}")
     
     progress(0.9, desc="Saving outputs...")
@@ -310,7 +430,22 @@ def generate_3d_ui(prompt, progress=gr.Progress()):
         with open(glb_path, 'wb') as f:
             f.write(glb_bytes)
     
+    # Save logs
+    images_dict = {
+        "original": original_path,
+        "transparent": nobg_path,
+        "mask": mask_path
+    }
+    metadata = {
+        "type": "ui_generation",
+        "has_glb": glb_path is not None,
+        "ply_size_bytes": len(ply_bytes),
+        "glb_size_bytes": len(glb_bytes) if glb_bytes else 0
+    }
+    save_generation_log(prompt, images_dict, metadata)
+    
     progress(1.0, desc="Done!")
+    logger.info(f"✓ Generation complete!")
     
     return (
         original_path,
@@ -326,6 +461,9 @@ def generate_3d_ui(prompt, progress=gr.Progress()):
 
 def edit_3d_ui(edit_prompt, current_image, edit_count, progress=gr.Progress()):
     """UI wrapper for editing"""
+    logger.info(f"=== EDIT REQUEST #{edit_count + 1} ===")
+    logger.info(f"Edit: {edit_prompt}")
+    
     if current_image is None:
         raise gr.Error("No image to edit. Generate a 3D model first!")
     
@@ -357,8 +495,11 @@ def edit_3d_ui(edit_prompt, current_image, edit_count, progress=gr.Progress()):
         
         if edited_image is None:
             raise gr.Error("Edit failed")
+        
+        logger.info("✓ Edit applied")
             
     except Exception as e:
+        logger.error(f"✗ Edit failed: {e}")
         raise gr.Error(f"Edit failed: {e}")
     
     progress(0.3, desc="Creating new mask...")
@@ -369,6 +510,7 @@ def edit_3d_ui(edit_prompt, current_image, edit_count, progress=gr.Progress()):
     try:
         ply_bytes, glb_bytes = run_sam3d(edited_image, gray)
     except Exception as e:
+        logger.error(f"✗ SAM-3D failed: {e}")
         raise gr.Error(f"SAM-3D failed: {e}")
     
     progress(0.9, desc="Saving outputs...")
@@ -392,7 +534,21 @@ def edit_3d_ui(edit_prompt, current_image, edit_count, progress=gr.Progress()):
             f.write(glb_bytes)
     
     new_edit_count = edit_count + 1
+    
+    # Save logs
+    images_dict = {
+        "edited": nobg_path,
+        "mask": mask_path
+    }
+    metadata = {
+        "type": "ui_edit",
+        "edit_number": new_edit_count,
+        "has_glb": glb_path is not None
+    }
+    save_generation_log(edit_prompt, images_dict, metadata)
+    
     progress(1.0, desc=f"Edit #{new_edit_count} complete!")
+    logger.info(f"✓ Edit #{new_edit_count} complete!")
     
     return (
         nobg_path,
@@ -529,4 +685,5 @@ demo = gr.TabbedInterface(
 )
 
 if __name__ == "__main__":
+    logger.info("=== Starting Text-to-3D MCP Server ===")
     demo.launch(mcp_server=True)