Update app.py

app.py

@@ -1,10 +1,33 @@
+import gradio as gr
+import torch
+from transformers import pipeline, set_seed
 from diffusers import DiffusionPipeline
 import tempfile
-
-# Load video generation model (you can switch to others)
-video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16)
-video_pipeline = video_pipeline.to("cuda")
-
+import imageio
+
+# ---------- Models ----------
+AVAILABLE_MODELS = {
+    "GPT-2 (small, fast)": "gpt2",
+    "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
+    "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
+}
+
+set_seed(42)
+text_model_cache = {}
+image_generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v1-4")
+
+# Try loading video model safely
+try:
+    video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
+    video_enabled = True
+except Exception as e:
+    video_pipeline = None
+    video_enabled = False
+    print(f"[Video model not loaded]: {e}")
+
+chat_memory = {}
+
+# ---------- Core Function ----------
 def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
@@ -13,42 +36,52 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         chat_memory[session_id] = []
         return "🧠 Codette signing off... Session reset.", None, None
 
-    generator = load_text_model(model_name)
+    # Load and run text model
+    if model_name not in text_model_cache:
+        text_model_cache[model_name] = pipeline("text-generation", model=AVAILABLE_MODELS[model_name])
+    generator = text_model_cache[model_name]
     response = generator(prompt, max_length=100, num_return_sequences=1, do_sample=True)[0]['generated_text'].strip()
 
     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:])
 
-    image = None
-    video = None
-
-    if generate_image:
-        image = image_generator(prompt)[0]['image']
-
-    if generate_video:
-        # Generate video from prompt
-        video_frames = video_pipeline(prompt, num_inference_steps=50).frames
-        # Save frames as video
-        temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-        import imageio
-        imageio.mimsave(temp_video_path, video_frames, fps=8)
-        video = temp_video_path
-
-    return chat_log, image, video
-
-
-
-    generate_video_toggle = gr.Checkbox(label="Also generate video?", value=False)
-output_video = gr.Video(label="AI-Generated Video")
-
-user_input.submit(
-    fn=codette_terminal,
-    inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
-    outputs=[output_text, output_image, output_video]
-)
-
-
-
-
-
+    # Image generation
+    image = image_generator(prompt)[0]['image'] if generate_image else None
+
+    # Video generation (if enabled and requested)
+    video_path = None
+    if generate_video and video_enabled:
+        try:
+            video_frames = video_pipeline(prompt, num_inference_steps=50).frames
+            temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+            imageio.mimsave(temp_video_path, video_frames, fps=8)
+            video_path = temp_video_path
+        except Exception as ve:
+            chat_log += f"\n[Video generation error]: {ve}"
+
+    return chat_log, image, video_path
+
+# ---------- UI ----------
+with gr.Blocks(title="Codette Terminal – AI Text + Image + Video") as demo:
+    gr.Markdown("## 🧬 Codette Terminal (Text + Image + Video, Hugging Face Edition)")
+    gr.Markdown("Choose a model, type your prompt, and optionally generate an image or video.\nType `'exit'` to reset 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 Language Model")
+    generate_image_toggle = gr.Checkbox(label="Also generate image?", value=False)
+    generate_video_toggle = gr.Checkbox(label="Also generate video?", value=False, interactive=video_enabled)
+
+    user_input = gr.Textbox(label="Your Prompt", placeholder="e.g. A dragon flying over Tokyo", lines=1)
+    output_text = gr.Textbox(label="Codette Output", lines=15, interactive=False)
+    output_image = gr.Image(label="Generated Image")
+    output_video = gr.Video(label="Generated Video")
+
+    user_input.submit(
+        fn=codette_terminal,
+        inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
+        outputs=[output_text, output_image, output_video]
+    )
+
+if __name__ == "__main__":
+    demo.launch()