Update app.py

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
-import numpy as np
 import tempfile
 import imageio
 import torch
 from transformers import pipeline
 from diffusers import DiffusionPipeline
 
-# ---------- Settings ----------
+# ---------- Configuration ----------
 AVAILABLE_MODELS = {
     "GPT-2 (small, fast)": "gpt2",
     "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
@@ -17,59 +16,87 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 text_model_cache = {}
 chat_memory = {}
 
-# ---------- Load Optional Models ----------
+# ---------- Load Image Generator ----------
 try:
-    image_generator = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+    image_generator = DiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
     image_generator.to(device)
     image_enabled = True
 except Exception as e:
-    print(f"[Image Load Error]: {e}")
+    print(f"[Image Model Load Error]: {e}")
     image_generator = None
     image_enabled = False
 
+# ---------- Load Video Generator ----------
 try:
-    video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+    video_pipeline = DiffusionPipeline.from_pretrained(
+        "damo-vilab/text-to-video-ms-1.7b",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
     video_pipeline.to(device)
     video_enabled = True
 except Exception as e:
-    print(f"[Video Load Error]: {e}")
+    print(f"[Video Model Load Error]: {e}")
     video_pipeline = None
     video_enabled = False
 
-# ---------- Core Logic ----------
+# ---------- Streamed Response Generator ----------
 def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
 
     if prompt.lower() in ["exit", "quit"]:
         chat_memory[session_id] = []
-        return "🧠 Codette signing off... Session reset.", None, None
+        yield "🧠 Codette signing off... Session reset.", None, None
+        return
 
-    # Load text model if not cached
+    # Load text model if not already loaded
     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)
+            text_model_cache[model_name] = pipeline(
+                "text-generation",
+                model=AVAILABLE_MODELS[model_name],
+                device=0 if device == "cuda" else -1
+            )
         except Exception as e:
-            return f"[Text model error]: {e}", None, None
+            yield f"[Text model 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)
-        response = output[0]['generated_text'].strip()
+        output = generator(prompt, max_length=100, do_sample=True, num_return_sequences=1)[0]['generated_text'].strip()
     except Exception as e:
-        response = f"[Text generation error]: {e}"
+        yield f"[Text generation error]: {e}", None, None
+        return
 
+    # Stream the output character by character
+    response_so_far = ""
+    for char in output:
+        response_so_far += char
+        temp_log = chat_memory[session_id][:]
+        temp_log.append(f"🖋️ You > {prompt}")
+        temp_log.append(f"🧠 Codette > {response_so_far}")
+        yield "\n".join(temp_log[-10:]), None, None
+        import time
+        time.sleep(0.01)
+
+    # Finalize chat memory
     chat_memory[session_id].append(f"🖋️ You > {prompt}")
-    chat_memory[session_id].append(f"🧠 Codette > {response}")
-    chat_log = "\n".join(chat_memory[session_id][-10:])
+    chat_memory[session_id].append(f"🧠 Codette > {output}")
 
+    # Image Generation
     img = None
     if generate_image and image_enabled:
         try:
             img = image_generator(prompt).images[0]
         except Exception as e:
-            chat_log += f"\n[Image error]: {e}"
+            response_so_far += f"\n[Image error]: {e}"
 
+    # Video Generation
     vid = None
     if generate_video and video_enabled:
         try:
@@ -78,13 +105,15 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
             imageio.mimsave(temp_video_path, video_frames, fps=8)
             vid = temp_video_path
         except Exception as e:
-            chat_log += f"\n[Video error]: {e}"
+            response_so_far += f"\n[Video error]: {e}"
+
+    yield "\n".join(chat_memory[session_id][-10:]), img, vid
 
-    return chat_log, img, vid
+# ---------- Gradio UI ----------
+with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
+    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video)")
+    gr.Markdown("Type a prompt and select your model. Enable image or video generation if you like. Type `'exit'` to reset.")
 
-# ---------- Gradio App ----------
-with gr.Blocks(title="Codette Terminal (Hugging Face Edition)") as demo:
-    gr.Markdown("## 🧬 Codette Terminal\nA text + image + video AI powered by Hugging Face + Gradio. Type `'exit'` to reset the session.")
     session_id = gr.Textbox(value="session_default", visible=False)
     model_dropdown = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Choose a Language Model")
     generate_image_toggle = gr.Checkbox(label="Also generate image?", value=False, interactive=image_enabled)
@@ -97,8 +126,12 @@ with gr.Blocks(title="Codette Terminal (Hugging Face Edition)") as demo:
     user_input.submit(
         codette_terminal,
         inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
-        outputs=[output_text, output_image, output_video]
+        outputs=[output_text, output_image, output_video],
+        concurrency_limit=1,
+        queue=True,
+        show_progress=True
     )
 
+# ---------- Launch ----------
 if __name__ == "__main__":
     demo.launch()