Update app.py

app.py

@@ -1,71 +1,57 @@
-from flask import Flask, request, jsonify, send_from_directory
-import requests
-import io
-import random
-import os
-from PIL import Image
-from datetime import datetime
+from flask import Flask, request, jsonify
+from huggingface_hub import InferenceClient
 
 app = Flask(__name__)
+app.config["DEBUG"] = True  # Enable for debugging
 
-API_URL = "https://api-inference.huggingface.co/models/openskyml/dalle-3-xl"
-API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
+# Load model client
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
-TEMP_DIR = "temp"
+# Function for text generation with enhanced prompt formatting
+def generate(
+    prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
 
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
 
-def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None):
-    payload = {
-        "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed if seed is not None else random.randint(-1, 2147483647)
-    }
+    # Enhanced prompt formatting for better context
+    formatted_prompt = f"{system_prompt}\n" + ''.join(f"{user_prompt} ||| {bot_response}\n" for user_prompt, bot_response in history) + f"\n{prompt}"
 
-    image_bytes = requests.post(API_URL, headers=headers, json=payload).content
-    image = Image.open(io.BytesIO(image_bytes))
-    return image
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
 
-
-def save_image(image, prompt):
-    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-    filename = f"{prompt}-{random.randint(1, 100000)}-{timestamp}.png"
-    filepath = os.path.join(TEMP_DIR, filename)
-    image.save(filepath, format='PNG')
-    return filename
+    for response in stream:
+        output += response.token.text
+    return output
 
 
 @app.route("/generate", methods=["POST"])
-def generate():
-    try:
-        data = request.get_json()
-        prompt = data["prompt"]
-        negative_prompt = data.get("negative_prompt", "")
-        is_negative = True if negative_prompt else False
-
-        image = query(prompt, is_negative=is_negative)
-        filename = save_image(image, prompt)
-
-        response = {
-            "success": True,
-            "image_filename": filename
-        }
-    except Exception as e:
-        response = {
-            "success": False,
-            "error": str(e)
-        }
-
-    return jsonify(response)
-
-
-@app.route("/temp/<filename>")
-def show_image(filename):
-    return send_from_directory(TEMP_DIR, filename)
-
-
+def generate_text():
+    data = request.json
+    prompt = data.get("prompt")
+    history = data.get("history", [])
+    system_prompt = data.get("system_prompt")
+    temperature = data.get("temperature", 0.9)
+    max_new_tokens = data.get("max_new_tokens", 256)
+    top_p = data.get("top_p", 0.95)
+    repetition_penalty = data.get("repetition_penalty", 1.0)
+
+    response = generate(
+        prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty
+    )
+
+    return jsonify({"response": response})
 if __name__ == "__main__":
     if not os.path.exists(TEMP_DIR):
         os.makedirs(TEMP_DIR)