Update app.py

app.py

@@ -3,19 +3,17 @@ import tempfile
 import imageio
 import torch
 import time
+import os
+import openai
 from transformers import pipeline
 from diffusers import DiffusionPipeline
 
+# ---------- Secret Key Setup ----------
+openai.api_key = os.getenv("OPENAI_KEY")  # ✅ Hugging Face Space Secret
 
-
-user_input.submit(
-    codette_terminal_limited,  # <== use your chosen function name here
-    inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id, batch_size_slider, video_steps_slider, fps_slider],
-    outputs=[output_text, output_image, output_video]
-)
-
-# ---------- Configuration ----------
+# ---------- Config ----------
 AVAILABLE_MODELS = {
+    "Codette Fine-Tuned (v9)": "ft:gpt-4.1-2025-04-14:raiffs-bits:codette-final:BO907H7Z",
     "GPT-2 (small, fast)": "gpt2",
     "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
     "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
@@ -24,6 +22,9 @@ AVAILABLE_MODELS = {
 device = "cuda" if torch.cuda.is_available() else "cpu"
 text_model_cache = {}
 chat_memory = {}
+MAX_PROMPTS_PER_SESSION = 5
+THROTTLE_SECONDS = 30
+last_usage_time = {}
 
 # ---------- Load Image Generator ----------
 try:
@@ -53,8 +54,8 @@ except Exception as e:
     video_pipeline = None
     video_enabled = False
 
-# ---------- Streamed Response Generator ----------
-def codette_terminal(prompt, model_name, generate_image, generate_video, session_id, batch_size, video_steps, fps):
+# ---------- Rate-Limited Terminal ----------
+def codette_terminal_limited(prompt, model_name, generate_image, generate_video, session_id, batch_size, video_steps, fps):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
 
@@ -63,28 +64,47 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         yield "🧠 Codette signing off... Session reset.", None, None
         return
 
-    # Load text model if not already loaded
-    if model_name not in text_model_cache:
+    if model_name == "Codette Fine-Tuned (v9)":
+        count = sum(1 for line in chat_memory[session_id] if line.startswith("🖋️ You >"))
+        if count >= MAX_PROMPTS_PER_SESSION:
+            yield "[🛑 Limit] Max 5 prompts per session.", None, None
+            return
+        now = time.time()
+        if now - last_usage_time.get(session_id, 0) < THROTTLE_SECONDS:
+            wait = int(THROTTLE_SECONDS - (now - last_usage_time[session_id]))
+            yield f"[⏳ Wait] Try again in {wait} sec.", None, None
+            return
+        last_usage_time[session_id] = now
+
+    if model_name == "Codette Fine-Tuned (v9)":
         try:
-            text_model_cache[model_name] = pipeline(
-                "text-generation",
+            response = openai.ChatCompletion.create(
                 model=AVAILABLE_MODELS[model_name],
-                device=0 if device == "cuda" else -1
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.7,
+                max_tokens=256
             )
+            output = response.choices[0].message.content.strip()
         except Exception as e:
-            yield f"[Text model error]: {e}", None, None
+            yield f"[OpenAI error]: {e}", None, None
+            return
+    else:
+        if model_name not in text_model_cache:
+            try:
+                text_model_cache[model_name] = pipeline(
+                    "text-generation",
+                    model=AVAILABLE_MODELS[model_name],
+                    device=0 if device == "cuda" else -1
+                )
+            except Exception as e:
+                yield f"[Text model error]: {e}", None, None
+                return
+        try:
+            output = text_model_cache[model_name](prompt, max_length=100, do_sample=True, num_return_sequences=1)[0]['generated_text'].strip()
+        except Exception as e:
+            yield f"[Generation error]: {e}", None, None
             return
 
-    generator = text_model_cache[model_name]
-
-    # Generate response
-    try:
-        output = generator(prompt, max_length=100, do_sample=True, num_return_sequences=1)[0]['generated_text'].strip()
-    except Exception as e:
-        yield f"[Text generation error]: {e}", None, None
-        return
-
-    # Stream the output
     response_so_far = ""
     for char in output:
         response_so_far += char
@@ -94,11 +114,10 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         yield "\n".join(temp_log[-10:]), None, None
         time.sleep(0.01)
 
-    # Finalize chat memory
     chat_memory[session_id].append(f"🖋️ You > {prompt}")
     chat_memory[session_id].append(f"🧠 Codette > {output}")
 
-    imgs = None
+    imgs, vid = None, None
     if generate_image and image_enabled:
         try:
             result = image_generator(prompt, num_images_per_prompt=batch_size)
@@ -106,7 +125,6 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         except Exception as e:
             response_so_far += f"\n[Image error]: {e}"
 
-    vid = None
     if generate_video and video_enabled:
         try:
             result = video_pipeline(prompt, num_inference_steps=video_steps)
@@ -121,8 +139,8 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
 
 # ---------- Gradio UI ----------
 with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
-    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video + Batch + NSFW OK)")
-    gr.Markdown("Type a prompt, select your model, and configure generation options. Type `'exit'` to reset.")
+    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video + Rate-Limited Access)")
+    gr.Markdown("Type a prompt, select a model, and generate responses, images, or video. Type `'exit'` to reset.")
 
     with gr.Row():
         session_id = gr.Textbox(value="session_default", visible=False)
@@ -143,7 +161,7 @@ with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
     output_video = gr.Video(label="Generated Video")
 
     user_input.submit(
-        codette_terminal,
+        codette_terminal_limited,
         inputs=[
             user_input, model_dropdown, generate_image_toggle, generate_video_toggle,
             session_id, batch_size_slider, video_steps_slider, fps_slider
@@ -153,4 +171,4 @@ with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
 
 # ---------- Launch ----------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(mcp_server=True)  # ✅ tool-ready