Update app.py

app.py

@@ -8,7 +8,7 @@ from openai import OpenAI
 from transformers import pipeline
 from diffusers import DiffusionPipeline
 
-# ---------- Load OpenAI Key from HF Secrets ----------
+# ---------- Load OpenAI Key from Hugging Face Secret ----------
 client = OpenAI(api_key=os.getenv("OPENAI_KEY"))
 
 # ---------- Configuration ----------
@@ -55,8 +55,9 @@ except Exception as e:
     video_pipeline = None
     video_enabled = False
 
-# ---------- Main Terminal with Usage Limits ----------
-def codette_terminal_limited(prompt, model_name, generate_image, generate_video, session_id, batch_size, video_steps, fps):
+# ---------- Main Terminal with Rate 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] = []
 
@@ -87,7 +88,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 v1 error]: {e}", None, None
+            yield f"[OpenAI error]: {e}", None, None
             return
     else:
         if model_name not in text_model_cache:
@@ -108,7 +109,7 @@ def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
             yield f"[Generation error]: {e}", None, None
             return
 
-    # Stream the output
+    # Stream text output
     response_so_far = ""
     for char in output:
         response_so_far += char
@@ -121,8 +122,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}")
 
-    # Image Generation
-    imgs = None
+    imgs, vid = None, None
+
     if generate_image and image_enabled:
         try:
             result = image_generator(prompt, num_images_per_prompt=batch_size)
@@ -130,8 +131,6 @@ 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)
@@ -147,7 +146,7 @@ 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 + Fine-Tuned AI)")
-    gr.Markdown("Type a prompt, choose a model, and generate intelligent responses. Type `'exit'` to reset the session.")
+    gr.Markdown("Type a prompt, choose a model, and generate responses. Type `'exit'` to reset the session.")
 
     with gr.Row():
         session_id = gr.Textbox(value="session_default", visible=False)
@@ -162,11 +161,29 @@ 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
+    with gr.Row():
+        user_input = gr.Textbox(
+            label="Your Prompt",
+            placeholder="e.g. A robot dreaming on Mars",
+            lines=1
+        )
+
+    with gr.Row():
+        output_text = gr.Textbox(label="Codette Output", lines=15, interactive=False)
+
+    with gr.Row():
+        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]
     )
-    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")
+
+# ---------- Launch ----------
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)