Spaces:
Runtime error
Runtime error
File size: 1,756 Bytes
0a74686 48056a7 0a74686 9a14904 48056a7 0a74686 1087492 9a14904 583e56c 81914fc 1087492 0a74686 9a14904 0a74686 9a14904 583e56c 0a74686 583e56c 0a74686 9a14904 583e56c 9a14904 0a74686 1087492 48056a7 583e56c 9a14904 583e56c 1087492 9a14904 3831488 583e56c 0a74686 583e56c 0a74686 583e56c 0a74686 583e56c 7949d53 0a74686 7949d53 1087492 48056a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import io
import base64
import torch
from flask import Flask, request, jsonify, send_file
from diffusers import DiffusionPipeline
from PIL import Image
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = Flask(__name__)
# Load the model once at startup (on CPU)
try:
logger.info("Loading Zero123Plus pipeline...")
pipe = DiffusionPipeline.from_pretrained(
"sudo-ai/zero123plus-v1.2",
torch_dtype=torch.float32, # CPU needs float32
)
pipe.to("cpu")
logger.info("=== Application Startup at CPU mode =====")
except Exception as e:
logger.error(f"Error loading model: {e}")
pipe = None
def pil_to_base64(image):
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode("utf-8")
@app.route("/")
def home():
return "Zero123Plus CPU API is running!"
@app.route("/generate", methods=["POST"])
def generate():
if pipe is None:
return jsonify({"error": "Model not loaded"}), 500
try:
data = request.get_json()
image_data = data.get("image")
if not image_data:
return jsonify({"error": "No image provided"}), 400
if image_data.startswith("data:image"):
image_data = image_data.split(",")[1]
image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
result = pipe(image)
output_image = result.images[0]
return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
except Exception as e:
logger.error(f"Error generating image: {e}")
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860) |