3D_Model_AI

Sleeping

App Files Files Community

everydaycats commited on Nov 22, 2025

Commit

fe1a826

verified ·

1 Parent(s): 84a8f0a

Update main.py

Browse files

Files changed (1) hide show

main.py +157 -61

main.py CHANGED Viewed

@@ -1,25 +1,80 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException, Query, Form
-from fastapi.responses import FileResponse, HTMLResponse
 import os
-import shutil
 import uuid
 import requests
 from typing import Optional
-from app.utils import run_inference
 from huggingface_hub import login
 hf_token = os.environ.get("HF_TOKEN")
-print(hf_token)
-login(token=hf_token)
 app = FastAPI(title="Stable Fast 3D API")
-# Embedded HTML page on root
 @app.get("/", response_class=HTMLResponse)
-async def root():
     html = """
     <!doctype html>
     <html>
@@ -27,7 +82,6 @@ async def root():
         <meta charset="utf-8" />
         <title>Stable Fast 3D API — Demo</title>
         <meta name="viewport" content="width=device-width,initial-scale=1" />
-        <!-- model-viewer for GLB preview -->
         <script type="module" src="https://unpkg.com/@google/model-viewer/dist/model-viewer.min.js"></script>
         <style>
           body { font-family: Inter, system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial; padding: 24px; background:#f7f8fb; color:#111; }
@@ -63,7 +117,6 @@ async def root():
           <div class="status" id="status">Status: idle</div>
           <a id="downloadLink" class="download-link" style="display:none" target="_blank">Download model (click if not auto-downloaded)</a>
-          <!-- Model preview -->
           <model-viewer id="preview" camera-controls auto-rotate environment-image="neutral" style="display:none;"></model-viewer>
         </div>
@@ -108,7 +161,6 @@ async def root():
               if (hasFile) {
                 formData.append('image', fileInput.files[0]);
               } else {
-                // we send the image_url as form field; server accepts this in POST
                 formData.append('image_url', urlInput.value.trim());
               }
@@ -129,10 +181,8 @@ async def root():
               const contentDisposition = resp.headers.get('content-disposition') || '';
               const suggestedName = (contentDisposition.match(/filename=([^;]+)/) || [])[1] || 'model.glb';
-              // Create object URL
               const url = URL.createObjectURL(blob);
-              // Auto-download
               const a = document.createElement('a');
               a.href = url;
               a.download = suggestedName.replace(/"/g, '');
@@ -140,13 +190,11 @@ async def root():
               a.click();
               a.remove();
-              // Show download link in case auto-download blocked
               downloadLink.href = url;
               downloadLink.download = a.download;
               downloadLink.textContent = 'Download ' + a.download;
               downloadLink.style.display = 'inline-block';
-              // Try to preview with model-viewer
               try {
                 preview.src = url;
                 preview.style.display = 'block';
@@ -155,8 +203,7 @@ async def root():
               }
               status.textContent = 'Status: Done — model generated.';
-              // Release URL after some time (keeping it available for a bit)
-              setTimeout(() => URL.revokeObjectURL(url), 5 * 60 * 1000); // revoke after 5 min
             } catch (err) {
               console.error(err);
@@ -172,76 +219,125 @@ async def root():
     return HTMLResponse(content=html, status_code=200)
-# Accept both an uploaded file OR an image_url via POST form
-@app.post("/generate-3d/")
-async def generate_3d_model_upload(
     image: Optional[UploadFile] = File(None),
     image_url: Optional[str] = Form(None),
 ):
-    """Generate 3D model from uploaded image file or from provided image URL (POST)"""
-    return await process_image(image=image, image_url=image_url)
-# Keep the existing GET-based URL endpoint (for direct API use)
-@app.get("/generate-3d/")
-async def generate_3d_model_url(image_url: str = Query(..., description="URL of the image to convert to 3D")):
-    """Generate 3D model from image URL (GET)"""
-    return await process_image(image_url=image_url)
-async def process_image(image: Optional[UploadFile] = None, image_url: Optional[str] = None):
-    # Create unique ID for this request
-    temp_id = str(uuid.uuid4())
-    tmp_dir = "/app/tmp"
-    os.makedirs(tmp_dir, exist_ok=True)
-    input_path = os.path.join(tmp_dir, f"{temp_id}.png")
-    output_dir = os.path.join(tmp_dir, f"{temp_id}_output")
     os.makedirs(output_dir, exist_ok=True)
     try:
-        # Handle image from upload or URL
         if image is not None:
-            with open(input_path, "wb") as f:
-                shutil.copyfileobj(image.file, f)
         elif image_url:
-            response = requests.get(image_url, stream=True, timeout=30)
-            if response.status_code != 200:
-                raise HTTPException(status_code=400, detail="Could not download image from URL")
-            with open(input_path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
         else:
-            raise HTTPException(status_code=400, detail="Either image file or image URL must be provided")
-        # Run the inference (your existing function)
-        glb_path = run_inference(input_path, output_dir)
         if not glb_path or not os.path.exists(glb_path):
-            raise HTTPException(status_code=500, detail="Failed to generate 3D model")
-        # Return the GLB file for download
         return FileResponse(
             path=glb_path,
             media_type="model/gltf-binary",
-            filename="model.glb",
-            headers={"Content-Disposition": f"attachment; filename=model.glb"},
         )
     except HTTPException:
         raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
-    finally:
-        # Clean up temporary files
         try:
             if os.path.exists(input_path):
                 os.remove(input_path)
-        except Exception:
-            pass
-        # Optionally, clean up output_dir contents (you could keep for debugging)
-        try:
             if os.path.exists(output_dir):
-                shutil.rmtree(output_dir)
         except Exception:
-            pass

+# app.py
 import os
 import uuid
+import shutil
+import logging
 import requests
 from typing import Optional
+from fastapi import FastAPI, UploadFile, File, HTTPException, Query, Form, BackgroundTasks
+from fastapi.responses import FileResponse, HTMLResponse
 from huggingface_hub import login
+from app.utils import run_inference
+# --- Configuration / env ---
 hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    try:
+        login(token=hf_token)
+    except Exception:
+        # If login fails, continue (some deployments don't need it)
+        pass
+TMP_DIR = os.environ.get("TMP_DIR", "/app/tmp")
+os.makedirs(TMP_DIR, exist_ok=True)
+# Logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("stable-fast-3d-api")
 app = FastAPI(title="Stable Fast 3D API")
+# --- Utility helpers ---
+def _cleanup_paths(input_path: Optional[str], output_dir: Optional[str]) -> None:
+    """Remove the input file and output directory after response has been sent."""
+    try:
+        if input_path and os.path.exists(input_path):
+            try:
+                os.remove(input_path)
+                logger.info("Cleaned input file: %s", input_path)
+            except Exception as e:
+                logger.warning("Failed to remove input file %s: %s", input_path, e)
+    except Exception as e:
+        logger.exception("Unexpected error while cleaning input file: %s", e)
+    try:
+        if output_dir and os.path.exists(output_dir):
+            try:
+                shutil.rmtree(output_dir, ignore_errors=True)
+                logger.info("Removed output directory: %s", output_dir)
+            except Exception as e:
+                logger.warning("Failed to remove output dir %s: %s", output_dir, e)
+    except Exception as e:
+        logger.exception("Unexpected error while cleaning output directory: %s", e)
+def _save_upload_file(upload_file: UploadFile, dest_path: str) -> None:
+    """Save UploadFile to destination path (binary stream)."""
+    with open(dest_path, "wb") as f:
+        shutil.copyfileobj(upload_file.file, f)
+    upload_file.file.close()
+def _download_to_file(url: str, dest_path: str, timeout: int = 30) -> None:
+    """Download a URL to a file (streaming)."""
+    resp = requests.get(url, stream=True, timeout=timeout)
+    if resp.status_code != 200:
+        raise HTTPException(status_code=400, detail=f"Failed to download image: status {resp.status_code}")
+    with open(dest_path, "wb") as f:
+        for chunk in resp.iter_content(chunk_size=8192):
+            if not chunk:
+                continue
+            f.write(chunk)
+# --- Embedded UI root ---
 @app.get("/", response_class=HTMLResponse)
+async def root_ui():
     html = """
     <!doctype html>
     <html>
         <meta charset="utf-8" />
         <title>Stable Fast 3D API — Demo</title>
         <meta name="viewport" content="width=device-width,initial-scale=1" />
         <script type="module" src="https://unpkg.com/@google/model-viewer/dist/model-viewer.min.js"></script>
         <style>
           body { font-family: Inter, system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial; padding: 24px; background:#f7f8fb; color:#111; }
           <div class="status" id="status">Status: idle</div>
           <a id="downloadLink" class="download-link" style="display:none" target="_blank">Download model (click if not auto-downloaded)</a>
           <model-viewer id="preview" camera-controls auto-rotate environment-image="neutral" style="display:none;"></model-viewer>
         </div>
               if (hasFile) {
                 formData.append('image', fileInput.files[0]);
               } else {
                 formData.append('image_url', urlInput.value.trim());
               }
               const contentDisposition = resp.headers.get('content-disposition') || '';
               const suggestedName = (contentDisposition.match(/filename=([^;]+)/) || [])[1] || 'model.glb';
               const url = URL.createObjectURL(blob);
               const a = document.createElement('a');
               a.href = url;
               a.download = suggestedName.replace(/"/g, '');
               a.click();
               a.remove();
               downloadLink.href = url;
               downloadLink.download = a.download;
               downloadLink.textContent = 'Download ' + a.download;
               downloadLink.style.display = 'inline-block';
               try {
                 preview.src = url;
                 preview.style.display = 'block';
               }
               status.textContent = 'Status: Done — model generated.';
+              setTimeout(() => URL.revokeObjectURL(url), 5 * 60 * 1000);
             } catch (err) {
               console.error(err);
     return HTMLResponse(content=html, status_code=200)
+# --- API endpoints: POST and GET ---
+@app.post("/generate-3d/", response_class=FileResponse)
+async def generate_3d_model_post(
     image: Optional[UploadFile] = File(None),
     image_url: Optional[str] = Form(None),
+    background_tasks: BackgroundTasks = None,
 ):
+    """
+    Accept either an uploaded image file (multipart/form-data) OR an image_url form field.
+    The final GLB will be returned as a FileResponse; cleanup is scheduled after response.
+    """
+    return await _process_image_and_respond(image=image, image_url=image_url, background_tasks=background_tasks)
+@app.get("/generate-3d/", response_class=FileResponse)
+async def generate_3d_model_get(
+    image_url: str = Query(..., description="URL of the image to convert to 3D"),
+    background_tasks: BackgroundTasks = None,
+):
+    return await _process_image_and_respond(image=None, image_url=image_url, background_tasks=background_tasks)
+# --- Core processing function ---
+async def _process_image_and_respond(
+    image: Optional[UploadFile] = None,
+    image_url: Optional[str] = None,
+    background_tasks: BackgroundTasks = None,
+):
+    """Handles request, writes input file, runs inference, returns FileResponse, schedules cleanup."""
+    request_id = str(uuid.uuid4())
+    input_path = os.path.join(TMP_DIR, f"{request_id}.png")
+    output_dir = os.path.join(TMP_DIR, f"{request_id}_output")
     os.makedirs(output_dir, exist_ok=True)
+    logger.info("New request %s: writing input & running inference", request_id)
     try:
+        # Save input image
         if image is not None:
+            logger.info("Saving uploaded file to %s", input_path)
+            _save_upload_file(image, input_path)
         elif image_url:
+            logger.info("Downloading image URL to %s", input_path)
+            try:
+                _download_to_file(image_url, input_path, timeout=30)
+            except HTTPException:
+                raise
+            except Exception as e:
+                logger.exception("Downloading image failed: %s", e)
+                raise HTTPException(status_code=400, detail=f"Failed to download image: {e}")
         else:
+            raise HTTPException(status_code=400, detail="Either image file or image_url must be provided")
+        # Run your heavy inference / external process. It must return a glb path.
+        logger.info("Running inference on image: %s", input_path)
+        try:
+            glb_path = run_inference(input_path, output_dir)
+        except Exception as e:
+            logger.exception("run_inference failed: %s", e)
+            # debug listing
+            listing = []
+            try:
+                for root, dirs, files in os.walk(output_dir):
+                    for fn in files:
+                        listing.append(os.path.join(root, fn))
+                logger.info("Output directory listing after failure: %s", listing)
+            except Exception as e2:
+                logger.exception("Failed to list output dir: %s", e2)
+            raise HTTPException(status_code=500, detail=f"run_inference failed: {e}")
+        # glb_path could be None or incorrect; verify exists
         if not glb_path or not os.path.exists(glb_path):
+            # If run_inference returns directory instead, try to locate a GLB inside
+            found = None
+            for root, dirs, files in os.walk(output_dir):
+                for fn in files:
+                    if fn.lower().endswith(".glb"):
+                        found = os.path.join(root, fn)
+                        break
+                if found:
+                    break
+            if found:
+                glb_path = found
+                logger.info("Found GLB via fallback discovery: %s", glb_path)
+            else:
+                listing = []
+                for root, dirs, files in os.walk(output_dir):
+                    for fn in files:
+                        listing.append(os.path.join(root, fn))
+                logger.error("GLB not produced. output_dir listing: %s", listing)
+                raise HTTPException(status_code=500, detail="Failed to generate 3D model (GLB missing). Check server logs.")
+        # Schedule cleanup to run after response completes
+        if background_tasks is not None:
+            background_tasks.add_task(_cleanup_paths, input_path, output_dir)
+            logger.info("Scheduled cleanup for request %s", request_id)
+        else:
+            logger.info("No BackgroundTasks provided; output will not be deleted automatically for request %s", request_id)
+        # Return the GLB as FileResponse (Starlette will stat the file before sending)
+        filename = os.path.basename(glb_path)
+        logger.info("Returning GLB %s for request %s", glb_path, request_id)
         return FileResponse(
             path=glb_path,
             media_type="model/gltf-binary",
+            filename=filename,
         )
     except HTTPException:
+        # pass through HTTPException so FastAPI can send correct response
         raise
     except Exception as e:
+        logger.exception("Unhandled exception during process: %s", e)
+        # Attempt to cleanup synchronously on error (best-effort)
         try:
             if os.path.exists(input_path):
                 os.remove(input_path)
             if os.path.exists(output_dir):
+                shutil.rmtree(output_dir, ignore_errors=True)
         except Exception:
+            pass
+        raise HTTPException(status_code=500, detail=f"Server error while processing image: {e}")