app.py

import os
import torch
from flask import Flask, request, jsonify, send_file
from pipeline import Zero123PlusPipeline  # from your local pipeline.py
from PIL import Image
from io import BytesIO

app = Flask(__name__)

# Load the model once at startup (on CPU)
print("Loading Zero123Plus pipeline on CPU...")
pipe = Zero123PlusPipeline.from_pretrained(
    "sudo-ai/zero123plus-v1.2",
    torch_dtype=torch.float32,
)
pipe.to("cpu")
pipe.enable_model_cpu_offload()
print("Model loaded.")

@app.route("/")
def home():
    return '''
        <h1>Zero123Plus Image to 3D Generator</h1>
        <form action="/generate" method="post" enctype="multipart/form-data">
            <p>Upload a single-view image:</p>
            <input type="file" name="image"><br><br>
            <input type="submit" value="Generate 3D View">
        </form>
    '''

@app.route("/generate", methods=["POST"])
def generate():
    if "image" not in request.files:
        return jsonify({"error": "No image uploaded"}), 400

    file = request.files["image"]
    image = Image.open(file.stream).convert("RGB")

    print("Generating 3D view...")
    result = pipe(image, num_inference_steps=50, guidance_scale=3.0)

    output = result.images[0]
    img_io = BytesIO()
    output.save(img_io, "PNG")
    img_io.seek(0)
    return send_file(img_io, mimetype="image/png")

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)