app.py

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
import base64

app = Flask(__name__)
CORS(app)

client = InferenceClient("briaai/RMBG-1.4")

@app.route("/")
def home():
    return {"message": "✅ Rembg Docker App Running"}

@app.route("/remove-bg", methods=["POST"])
def remove_bg():
    try:
        if "image" not in request.files:
            return jsonify({"error": "No image provided"}), 400

        # Read uploaded image
        image_file = request.files["image"]
        image = Image.open(image_file.stream)

        # Convert to bytes for the model
        buffer = BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        image_bytes = buffer.read()

        # Send image to Hugging Face model
        result = client.post(
            model="briaai/RMBG-1.4",
            data=image_bytes,
            headers={"Content-Type": "image/png"},
        )

        # The model returns processed image bytes
        if not result:
            return jsonify({"error": "Empty response from model"}), 500

        # Convert result to base64 string
        encoded_image = base64.b64encode(result).decode("utf-8")
        return jsonify({"image": encoded_image})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


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