Update app.py

app.py

@@ -1,8 +1,8 @@
 import io
 import base64
 import torch
-from flask import Flask, request, jsonify, send_file
-from diffusers import Zero123PlusPipeline
+from flask import Flask, request, jsonify
+from diffusers import StableDiffusionPipeline  # Placeholder; SF3D uses a custom setup
 from PIL import Image
 import logging
 
@@ -13,10 +13,10 @@ app = Flask(__name__)
 
 # Load the model once at startup (on CPU)
 try:
-    logger.info("Loading Zero123Plus pipeline...")
-    pipe = Zero123PlusPipeline.from_pretrained(
-        "sudo-ai/zero123plus-v1.2",
-        torch_dtype=torch.float32,  # CPU needs float32
+    logger.info("Loading Stable Fast 3D pipeline...")
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-fast-3d",
+        torch_dtype=torch.float32,
         cache_dir="/tmp/hf_home",
     )
     pipe.to("cpu")
@@ -32,7 +32,7 @@ def pil_to_base64(image):
 
 @app.route("/")
 def home():
-    return "Zero123Plus CPU API is running!"
+    return "Stable Fast 3D CPU API is running!"
 
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -52,14 +52,21 @@ def generate():
         image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
 
         logger.info("Processing image with pipeline...")
-        result = pipe(image)
-        output_image = result.images[0]
-        logger.info("Image processed successfully")
+        # SF3D-specific generation; adjust based on documentation
+        result = pipe(image)  # Placeholder; check SF3D repo for exact method
+        output_mesh = result.mesh  # Hypothetical output; may be .glb or .obj
 
-        return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
+        # Save mesh to temporary file and encode
+        output_path = "/tmp/output.glb"
+        output_mesh.save(output_path)
+        with open(output_path, "rb") as f:
+            mesh_data = base64.b64encode(f.read()).decode("utf-8")
+
+        logger.info("Mesh processed successfully")
+        return jsonify({"mesh": f"data:model/gltf-binary;base64,{mesh_data}"})
 
     except Exception as e:
-        logger.error(f"Error generating image: {e}", exc_info=True)
+        logger.error(f"Error generating mesh: {e}", exc_info=True)
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":