Update app.py

app.py

@@ -1,6 +1,7 @@
 from flask import Flask, request, jsonify, send_from_directory
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
+from google.generativeai import configure, GenerativeModel
 from PIL import Image
 from io import BytesIO
 from prodiapy import Prodia
@@ -254,45 +255,53 @@ async def generate_image():
         return jsonify({"status": "error", "error": "Internal Server Error"}), 500
 
 
-MODELS = {
-    'gemini-1.0-pro-001': 'gemini-1.0-pro-001',
-    'gemini-1.0-pro': 'gemini-1.0-pro',
-    'gemini-pro': 'gemini-pro',
-    'chat-bison-001': 'chat-bison-001'
+configure(api_key="AIzaSyCKc8gUTrb7KjhAwY6NwATaMWIDr5etKFk")
+
+# Set up the model
+generation_config = {
+    "temperature": 0.9,
+    "top_p": 1,
+    "top_k": 1,
+    "max_output_tokens": 2048,
 }
 
+safety_settings = [
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+    },
+]
+
+model = GenerativeModel(
+    model_name="gemini-1.0-pro-001",
+    generation_config=generation_config,
+    safety_settings=safety_settings
+)
+
 @app.route('/gemini', methods=['GET'])
 @limiter.limit("30 per minute")
 def gemini():
     prompt = request.args.get('prompt')
-    model = request.args.get('model', 'gemini-pro')
 
     if not prompt:
         return jsonify({'error': 'Prompt parameter is required'}), 400
 
-    if model not in MODELS:
-        return jsonify({'error': 'Invalid model'}), 400
-
-    url = "https://generativelanguage.googleapis.com/v1beta/models/{}/{}".format(MODELS[model], "generateContent" if model != 'chat-bison-001' else "generateMessage")
-    data = {
-        "contents": [
-            {
-                "parts": [
-                    {
-                        "text": prompt
-                    }
-                ]
-            }
-        ]
-    }
-
-    response = requests.post(url + "?key=AIzaSyCKc8gUTrb7KjhAwY6NwATaMWIDr5etKFk", json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        return jsonify(result)
-    else:
-        return jsonify({'error': 'Failed to generate content'}), 500
+    convo = model.start_chat(history=[prompt])
+    response = convo.last.text
+
+    return jsonify({'response': response})
     
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)