Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import os
 import torch
 import time
@@ -14,12 +15,9 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from trellis.pipelines import TrellisImageTo3DPipeline
+from diffusers import DiffusionPipeline
 
-# Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["ATTN_BACKEND"] = "native"  # Disable xformers/flash-attn
-os.environ["SPCONV_ALGO"] = "native"  # Optimize for CPU
 torch.set_default_device("cpu")
 torch.cuda.is_available = lambda: False
 torch.cuda.device_count = lambda: 0
@@ -27,18 +25,15 @@ torch.cuda.device_count = lambda: 0
 app = Flask(__name__)
 CORS(app)
 
-# Configure directories
 UPLOAD_FOLDER = '/tmp/uploads'
 RESULTS_FOLDER = '/tmp/results'
 CACHE_DIR = '/tmp/huggingface'
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
-# Create directories
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(RESULTS_FOLDER, exist_ok=True)
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Set Hugging Face cache
 os.environ['HF_HOME'] = CACHE_DIR
 os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
 os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
@@ -46,17 +41,13 @@ os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 
-# Job tracking
 processing_jobs = {}
-
-# Global model
-trellis_pipeline = None
+zero123_pipeline = None
 model_loaded = False
 model_loading = False
 
-# Configuration
-TIMEOUT_SECONDS = 360  # 6 minutes for TRELLIS
-MAX_DIMENSION = 256  # TRELLIS works with smaller images
+TIMEOUT_SECONDS = 300
+MAX_DIMENSION = 256
 
 class TimeoutError(Exception):
     pass
@@ -92,37 +83,33 @@ def process_with_timeout(function, args, timeout):
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-# Image preprocessing
 def preprocess_image(image_path):
     try:
         with Image.open(image_path) as img:
-            # Convert to RGB
             if img.mode == 'RGBA':
                 img = img.convert('RGB')
-            # Resize to 256x256
             img = img.resize((256, 256), Image.LANCZOS)
             return img
     except Exception as e:
         raise Exception(f"Error preprocessing image: {str(e)}")
 
 def load_model():
-    global trellis_pipeline, model_loaded, model_loading
+    global zero123_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return trellis_pipeline
+        return zero123_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return trellis_pipeline
+        return zero123_pipeline
     
     try:
         model_loading = True
-        print("Loading TRELLIS-image-large...")
+        print("Loading Zero123++...")
         
-        model_name = "JeffreyXiang/TRELLIS-image-large"
+        model_name = "sudo-ai/zero123plus-v1.2"
         
-        # Download model
         max_retries = 3
         retry_delay = 5
         for attempt in range(max_retries):
@@ -141,17 +128,16 @@ def load_model():
                 else:
                     raise
         
-        # Load TRELLIS pipeline
-        trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
+        zero123_pipeline = DiffusionPipeline.from_pretrained(
             model_name,
             cache_dir=CACHE_DIR,
             torch_dtype=torch.float32,
         )
-        trellis_pipeline.to("cpu")
+        zero123_pipeline.to("cpu")
         
         model_loaded = True
-        print("TRELLIS loaded successfully on CPU")
-        return trellis_pipeline
+        print("Zero123++ loaded successfully on CPU")
+        return zero123_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -162,18 +148,16 @@ def load_model():
 
 def generate_3d_model(image, detail_level):
     try:
-        # Parameters
-        num_steps = {'low': 50, 'medium': 75, 'high': 100}
+        num_steps = {'low': 30, 'medium': 50, 'high': 75}
         steps = num_steps[detail_level]
         
-        # Generate 3D model
         with torch.no_grad():
-            result = trellis_pipeline(image, num_inference_steps=steps, output_type="mesh")
+            result = zero123_pipeline(image, num_inference_steps=steps)
         
-        # Extract mesh
-        vertices = np.array(result.vertices)
-        faces = np.array(result.faces)
-        vertex_colors = np.array(result.vertex_colors) if result.vertex_colors is not None else None
+        mesh = result.meshes[0]
+        vertices = np.array(mesh.vertices)
+        faces = np.array(mesh.faces)
+        vertex_colors = np.array(mesh.vertex_colors) if mesh.vertex_colors is not None else None
         
         trimesh_mesh = trimesh.Trimesh(
             vertices=vertices,
@@ -181,7 +165,6 @@ def generate_3d_model(image, detail_level):
             vertex_colors=vertex_colors
         )
         
-        # Rotate for Unity Z-up
         trimesh_mesh.apply_transform(trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0]))
         
         return trimesh_mesh
@@ -192,7 +175,7 @@ def generate_3d_model(image, detail_level):
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "TRELLIS-image-large",
+        "model": "Zero123++",
         "device": "cpu"
     }), 200
 
@@ -303,7 +286,6 @@ def convert_image_to_3d():
                         
                 processing_jobs[job_id]['progress'] = 80
                 
-                # Export
                 file_path = os.path.join(output_dir, f"model.{output_format}")
                 mesh.export(file_path, file_type=output_format)
                 
@@ -369,7 +351,7 @@ def preview_model(job_id):
     if os.path.exists(file_path):
         if output_format == 'glb':
             return send_file(file_path, mimetype='model/gltf-binary')
-        else:  # OBJ
+        else:
             return send_file(file_path, mimetype='text/plain')
     
     return jsonify({"error": "Model file not found"}), 404
@@ -433,7 +415,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (TRELLIS-image-large)",
+        "message": "Image to 3D API (Zero123++)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -443,12 +425,13 @@ def index():
         ],
         "parameters": {
             "output_format": "glb or obj",
-            "detail_level": "low, medium, or high - controls inference steps"
+            "detail_level": "low, medium, or high"
         },
-        "description": "Creates 3D models from 2D images using TRELLIS-image-large. Use transparent PNGs for best results."
+        "description": "Creates 3D models from 2D images using Zero123++."
     }), 200
 
 if __name__ == '__main__':
     cleanup_old_jobs()
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)
+```