from flask import Flask, request, jsonify, send_file
import requests
import os
import base64
from io import BytesIO
from datetime import datetime
from flask_cors import CORS
import uuid

app = Flask(__name__)
CORS(app)  # 启用CORS支持

# NVIDIA API URL 和 API Key
NVIDIA_URL = "https://ai.api.nvidia.com/v1/genai/stabilityai/stable-diffusion-3-medium"
NVIDIA_API_KEY = os.environ.get('NVIDIA_API_KEY')

# 本地保存图片的目录
IMAGES_DIR = 'static/images'

@app.route('/deem/v1/images/generations', methods=['POST'])
def translate_request():
    data = request.json
    prompt = data.get('prompt')

    nvidia_payload = {
        "prompt": prompt,
        "cfg_scale": 5,
        "aspect_ratio": "16:9",
        "seed": 0,
        "steps": 50,
        "negative_prompt": ""
    }

    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Content-Type": "application/json"
    }

    # 请求 NVIDIA API
    response = requests.post(NVIDIA_URL, headers=headers, json=nvidia_payload)
    response.raise_for_status()
    response_body = response.json()

    # 获取base64编码的图片数据
    image_data = base64.b64decode(response_body['image'])

    # 生成唯一的文件名
    filename = f"{uuid.uuid4()}.jpg"
    file_path = os.path.join(IMAGES_DIR, filename)

    # 确保目录存在
    os.makedirs(IMAGES_DIR, exist_ok=True)

    # 将图片保存到本地
    with open(file_path, 'wb') as f:
        f.write(image_data)

    # 构造图片URL
    image_url = f"https://gitdeem-sd3.hf.space/static/images/{filename}"

    # 构造响应
    adapted_response = {
        "created": int(datetime.now().timestamp()),
        "data": [
            {
                "url": image_url
            }
        ]
    }

    return jsonify(adapted_response)

@app.route('/static/images/<filename>')
def serve_image(filename):
    return send_file(os.path.join(IMAGES_DIR, filename), mimetype='image/jpeg')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001)