feat: added model selector

2024-05-02 08:46:09 +02:00
parent 95f22c268b
commit a264e91f1f
2 changed files with 44 additions and 3 deletions
--- a/gen_service/main.py
+++ b/gen_service/main.py
@ -109,6 +109,7 @@ def _gen_text(prompt, use_zukijourney=False):
@app.route('/gen_image', methods=['GET'])
 def gen_image():
    prompt = request.args.get('prompt')
+    model = request.args.get('model', 'playground-v2.5')

    if not prompt:
        return "Error: No prompt provided", 400
@ -119,7 +120,7 @@ def gen_image():
    try:
        client = get_client()
        response = client.images.generate(
-            model="playground-v2.5",
+            model=model,
            prompt=prompt,
        )
        image_url = response.data[0].url