import os
os.system("pip install flask")

from flask import Flask, request, send_file, abort
from io import BytesIO
from PIL import Image
import random

app = Flask(__name__)

# モデルのロードをシミュレートするためのスタブ
def load_model(model_name):
    # モデルをロードするロジックをここに追加
    # ここでは、PILで空の画像を生成する簡易的な例を示します
    def model(prompt):
        img = Image.new('RGB', (256, 256), color = (73, 109, 137))
        return img
    return model

# モデルのロード
models_load = {f'model_{i}': load_model(f'model_{i}') for i in range(1, 7)}

@app.route('/generate_image', methods=['GET'])
def generate_image():
    prompt = request.args.get('prompt')
    model_str = request.args.get('model')
    
    if model_str not in models_load:
        abort(404, description="Model not found")
    
    model = models_load[model_str]
    noise = str(random.randint(0, 99999999999))
    image = model(f'{prompt} {noise}')
    
    # 画像をバイナリで保存し、レスポンスとして返す
    img_bytes = BytesIO()
    image.save(img_bytes, format='PNG')
    img_bytes.seek(0)
    
    return send_file(img_bytes, mimetype='image/png')

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