app.py

import os
import json
from flask import Flask, jsonify, request
from transformers import pipeline

# Create a Flask app
app = Flask(__name__)

# Initialize models at the start of the API
audio_model = None

def download_models():
    global audio_model
    print("Downloading models...")
    # Download and load the audio model
    audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
    print("Model downloaded and ready to use.")

# Download model when the server starts
download_models()

@app.route('/detect', methods=['POST'])
def detect_deepfake():
    data_type = request.form.get('type')  # "audio"
    audio_file = request.files.get('audio_file')
    folder_path = request.form.get('folder_path')

    # If a single audio file is provided
    if audio_file:
        try:
            # Save the uploaded file temporarily
            file_path = os.path.join("/tmp", audio_file.filename)
            audio_file.save(file_path)

            # Perform detection
            result = audio_model(file_path)
            result_dict = {item['label']: item['score'] for item in result}

            # Remove the temporary file
            os.remove(file_path)

            return jsonify({"message": "Detection completed", "results": result_dict}), 200

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

    # If a folder path is provided
    elif folder_path and os.path.isdir(folder_path):
        results = {}
        try:
            for file_name in os.listdir(folder_path):
                if file_name.endswith('.wav') or file_name.endswith('.mp3'):
                    file_path = os.path.join(folder_path, file_name)
                    result = audio_model(file_path)
                    results[file_name] = {item['label']: item['score'] for item in result}

            # Save results to a file
            with open('detection_results.json', 'w') as f:
                f.write(json.dumps(results))

            return jsonify({"message": "Detection completed", "results": results}), 200

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

    # Invalid request if neither audio file nor valid folder path is provided
    else:
        return jsonify({"error": "Invalid input. Provide an audio file or a valid folder path."}), 400

if __name__ == '__main__':
    # Run the Flask app
    app.run(host='0.0.0.0', port=7860)