Update app.py

app.py

@@ -1,42 +1,42 @@
-# This is a Gradio app that provides a text-based chat interface with optional image and video generation.
 import gradio as gr
 import numpy as np
 import tempfile
 import imageio
+import torch
+from transformers import pipeline
+from diffusers import DiffusionPipeline
 
-# ---------- Setup ----------
+# ---------- Settings ----------
 AVAILABLE_MODELS = {
     "GPT-2 (small, fast)": "gpt2",
     "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
     "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
 }
 
-# Load text-to-image model using diffusers (correct API)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+text_model_cache = {}
+chat_memory = {}
+
+# ---------- Load Optional Models ----------
 try:
-    from diffusers import DiffusionPipeline
-    image_generator = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    image_generator.to("cpu")
+    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 ImportError:
+except Exception as e:
+    print(f"[Image Load Error]: {e}")
     image_generator = None
     image_enabled = False
-    print("[Image model error]: diffusers library not found")
 
-# Load text-to-video model
 try:
-    from diffusers import DiffusionPipeline
-    video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
-    video_pipeline.to("cpu")
+    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 ImportError:
+except Exception as e:
+    print(f"[Video Load Error]: {e}")
     video_pipeline = None
     video_enabled = False
-    print("[Video model error]: diffusers library not found")
-
-text_model_cache = {}
-chat_memory = {}
 
-# ---------- Core Function ----------
+# ---------- Core Logic ----------
 def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
@@ -45,18 +45,24 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         chat_memory[session_id] = []
         return "🧠 Codette signing off... Session reset.", None, None
 
-    # Load the text generation model if it's not already in the cache
+    # Load text model if not cached
     if model_name not in text_model_cache:
-        text_model_cache[model_name] = gr.pipelines.TextGeneration(model=AVAILABLE_MODELS[model_name])
+        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:
+            return f"[Text model error]: {e}", None, None
+
     generator = text_model_cache[model_name]
-    response = generator(prompt, max_length=100, num_return_sequences=1, do_sample=True)[0]['generated_text'].strip()
+    try:
+        output = generator(prompt, max_length=100, do_sample=True, num_return_sequences=1)
+        response = output[0]['generated_text'].strip()
+    except Exception as e:
+        response = f"[Text generation error]: {e}"
 
-    # Update the chat log
     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:])
 
-    # Generate image if requested and image generation is enabled
     img = None
     if generate_image and image_enabled:
         try:
@@ -64,7 +70,6 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
         except Exception as e:
             chat_log += f"\n[Image error]: {e}"
 
-    # Generate video if requested and video generation is enabled
     vid = None
     if generate_video and video_enabled:
         try:
@@ -77,13 +82,11 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
 
     return chat_log, img, vid
 
-# ---------- Gradio UI ----------
-with gr.Blocks(title="Codette Terminal – Text + Image + Video") as demo:
-    gr.Markdown("## 🧬 Codette Terminal (Text + Image + Video, CPU-Friendly)")
-    gr.Markdown("Type a prompt and select your model. Optionally generate images or videos. Type `'exit'` to reset session.")
-
+# ---------- 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 Language Model")
+    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)
     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 robot dreaming on Mars", lines=1)
@@ -92,10 +95,10 @@ with gr.Blocks(title="Codette Terminal – Text + Image + Video") as demo:
     output_video = gr.Video(label="Generated Video")
 
     user_input.submit(
-        fn=codette_terminal,
+        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(show_error=True)
+    demo.launch()