Update app.py

app.py

@@ -1,53 +1,98 @@
-# This Gradio app allows users to interact with a chatbot that can generate text and images based on user prompts.
 import gradio as gr
-import numpy as np
-from transformers_js import pipeline  # Corrected import to use transformers_js instead of transformers_js_py
+import torch
+from transformers import pipeline, set_seed
+from diffusers import DiffusionPipeline
+import tempfile
+import imageio
 
-# Define the available models
+# ---------- Setup ----------
 AVAILABLE_MODELS = {
-    "GPT-2": "gpt2",
-    "DALL-E": "dalle-mini/dalle-mini-1.3B"
+    "GPT-2 (small, fast)": "gpt2",
+    "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
+    "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
 }
 
-# Initialize the text generation pipeline
-text_generator = pipeline("text-generation", model=AVAILABLE_MODELS["GPT-2"])
-
-# Initialize the image generation pipeline
-image_generator = pipeline("image-generation", model=AVAILABLE_MODELS["DALL-E"])
-
-# Function to generate text
-def generate_text(prompt, model):
-    np.random.seed(42)  # Set a seed for reproducibility
-    if model == "GPT-2":
-        return text_generator(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']
-    else:
-        return "Model not supported for text generation"
-
-# Function to generate images
-def generate_image(prompt, model):
-    if model == "DALL-E":
-        image = image_generator(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
-        return image
-    else:
-        return "Model not supported for image generation"
-
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Chatbot with Text and Image Generation")
-    
-    with gr.Tab("Text Generation"):
-        text_prompt = gr.Textbox(label="Enter your text prompt")
-        text_model = gr.Radio(choices=list(AVAILABLE_MODELS.keys()), label="Choose a model", value="GPT-2")
-        text_output = gr.Textbox(label="Generated Text")
-        text_button = gr.Button("Generate Text")
-        text_button.click(generate_text, inputs=[text_prompt, text_model], outputs=text_output)
-    
-    with gr.Tab("Image Generation"):
-        image_prompt = gr.Textbox(label="Enter your image prompt")
-        image_model = gr.Radio(choices=list(AVAILABLE_MODELS.keys()), label="Choose a model", value="DALL-E")
-        image_output = gr.Image(label="Generated Image")
-        image_button = gr.Button("Generate Image")
-        image_button.click(generate_image, inputs=[image_prompt, image_model], outputs=image_output)
-
-# Launch the interface
-demo.launch(show_error=True)
+set_seed(42)
+text_model_cache = {}
+
+# Load text-to-image model using diffusers (correct API)
+try:
+    image_generator = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    image_generator.to("cpu")
+    image_enabled = True
+except Exception as e:
+    image_generator = None
+    image_enabled = False
+    print(f"[Image model error]: {e}")
+
+# Load text-to-video model
+try:
+    video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
+    video_pipeline.to("cpu")
+    video_enabled = True
+except Exception as e:
+    video_pipeline = None
+    video_enabled = False
+    print(f"[Video model error]: {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] = []
+
+    if prompt.lower() in ["exit", "quit"]:
+        chat_memory[session_id] = []
+        return "🧠 Codette signing off... Session reset.", None, None
+
+    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:])
+
+    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}"
+
+    vid = 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)
+            vid = temp_video_path
+        except Exception as e:
+            chat_log += f"\n[Video error]: {e}"
+
+    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.")
+
+    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, 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)
+    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()