Update app.py

app.py

@@ -4,14 +4,14 @@ import imageio
 import torch
 import time
 import os
-import openai
+from openai import OpenAI
 from transformers import pipeline
 from diffusers import DiffusionPipeline
 
-# ---------- Secret Key Setup ----------
-openai.api_key = os.getenv("OPENAI_KEY")  # ✅ Hugging Face Space Secret
+# ---------- Load OpenAI Key from HF Secrets ----------
+client = OpenAI(api_key=os.getenv("OPENAI_KEY"))
 
-# ---------- Config ----------
+# ---------- Configuration ----------
 AVAILABLE_MODELS = {
     "Codette Fine-Tuned (v9)": "ft:gpt-4.1-2025-04-14:raiffs-bits:codette-final:BO907H7Z",
     "GPT-2 (small, fast)": "gpt2",
@@ -22,9 +22,10 @@ AVAILABLE_MODELS = {
 device = "cuda" if torch.cuda.is_available() else "cpu"
 text_model_cache = {}
 chat_memory = {}
+last_usage_time = {}
+
 MAX_PROMPTS_PER_SESSION = 5
 THROTTLE_SECONDS = 30
-last_usage_time = {}
 
 # ---------- Load Image Generator ----------
 try:
@@ -54,7 +55,7 @@ except Exception as e:
     video_pipeline = None
     video_enabled = False
 
-# ---------- Rate-Limited Terminal ----------
+# ---------- Main Terminal with Usage Limits ----------
 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] = []
@@ -78,7 +79,7 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
 
     if model_name == "Codette Fine-Tuned (v9)":
         try:
-            response = openai.ChatCompletion.create(
+            response = client.chat.completions.create(
                 model=AVAILABLE_MODELS[model_name],
                 messages=[{"role": "user", "content": prompt}],
                 temperature=0.7,
@@ -86,7 +87,7 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
             )
             output = response.choices[0].message.content.strip()
         except Exception as e:
-            yield f"[OpenAI error]: {e}", None, None
+            yield f"[OpenAI v1 error]: {e}", None, None
             return
     else:
         if model_name not in text_model_cache:
@@ -100,11 +101,14 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
                 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()
+            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
 
+    # Stream the output
     response_so_far = ""
     for char in output:
         response_so_far += char
@@ -117,7 +121,8 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
     chat_memory[session_id].append(f"🖋️ You > {prompt}")
     chat_memory[session_id].append(f"🧠 Codette > {output}")
 
-    imgs, vid = None, None
+    # Image Generation
+    imgs = None
     if generate_image and image_enabled:
         try:
             result = image_generator(prompt, num_images_per_prompt=batch_size)
@@ -125,6 +130,8 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
         except Exception as e:
             response_so_far += f"\n[Image error]: {e}"
 
+    # Video Generation
+    vid = None
     if generate_video and video_enabled:
         try:
             result = video_pipeline(prompt, num_inference_steps=video_steps)
@@ -139,8 +146,8 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
 
 # ---------- Gradio UI ----------
 with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
-    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.")
+    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video + Fine-Tuned AI)")
+    gr.Markdown("Type a prompt, choose a model, and generate intelligent responses. Type `'exit'` to reset the session.")
 
     with gr.Row():
         session_id = gr.Textbox(value="session_default", visible=False)
@@ -155,20 +162,4 @@ with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
         video_steps_slider = gr.Slider(label="Video Inference Steps", minimum=10, maximum=100, step=10, value=50)
         fps_slider = gr.Slider(label="Video FPS", minimum=4, maximum=24, step=2, value=8)
 
-    user_input = gr.Textbox(label="Your Prompt", placeholder="e.g. A robot dreaming on Mars", lines=1)
-    output_text = gr.Textbox(label="Codette Output", lines=15, interactive=False)
-    output_image = gr.Gallery(label="Generated Image(s)", columns=2)
-    output_video = gr.Video(label="Generated Video")
-
-    user_input.submit(
-        codette_terminal_limited,
-        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]
-    )
-
-# ---------- Launch ----------
-if __name__ == "__main__":
-    demo.launch(mcp_server=True)  # ✅ tool-ready
+    user_input = gr.Textbox(label="Your Prompt", placeholder="e._