import os
from flask import Flask, render_template, request, jsonify
from SadTalker import SadTalker

app = Flask(__name__)

# Initialize SadTalker with CPU
sadtalker = SadTalker(
    checkpoint_path="/app/SadTalker/checkpoints",
    config_path="/app/SadTalker/src/config",
    device="cpu"
)

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate():
    if 'image' not in request.files:
        return jsonify({"error": "No image uploaded"}), 400
    
    image = request.files['image']
    text = request.form.get('text', '')
    
    # Save files
    img_path = os.path.join('static/uploads', image.filename)
    audio_path = os.path.join('static/uploads', 'audio.wav')
    output_path = os.path.join('static/uploads', 'output.mp4')
    
    image.save(img_path)
    
    # Text-to-Speech (using gTTS)
    from gtts import gTTS
    tts = gTTS(text=text, lang='en')
    tts.save(audio_path)
    
    # Generate video (CPU optimized)
    sadtalker.generate(
        source_image=img_path,
        driven_audio=audio_path,
        result_dir='static/uploads',
        still=True,
        preprocess='crop',
        enhancer='none'  # Disable for CPU
    )
    
    return jsonify({
        "video": f"/static/uploads/{os.path.basename(output_path)}"
    })

if __name__ == '__main__':
    os.makedirs('static/uploads', exist_ok=True)
    app.run(host='0.0.0.0', port=7860)