Create app.py

app.py

@@ -0,0 +1,55 @@
+from flask import Flask, request, jsonify
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+app = Flask(__name__)
+
+# Load your model (can be any HF model)
+MODEL_NAME = "tiiuae/falcon-7b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+@app.route('/setToken', methods=['POST'])
+def set_token():
+    """
+    Main multimodal API endpoint.
+    Handles system + user prompts and returns generated response.
+    """
+    data = request.get_json(force=True)
+
+    system_prompt = data.get("system_prompt", "You are a helpful AI.")
+    user_input = data.get("user_input", "")
+    temperature = float(data.get("temperature", 0.7))
+    mode = data.get("mode", "text")
+
+    # Text mode (default)
+    if mode == "text":
+        full_prompt = f"{system_prompt}\nUser: {user_input}\nAI:"
+        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=temperature
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.split("AI:")[-1].strip()
+        return jsonify({
+            "model": MODEL_NAME,
+            "response": response,
+            "mode": "text"
+        })
+
+    # You can later add multimodal branches here:
+    # - "image" -> call image generation pipeline
+    # - "audio" -> call speech-to-text / text-to-speech
+    # - "embedding" -> return vector embeddings
+
+    return jsonify({"error": f"Unsupported mode: {mode}"}), 400
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860)