Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from transformers import pipeline, set_seed
 
-# Preload models into a dictionary for selection
+# Define available models
 AVAILABLE_MODELS = {
     "GPT-2 (small, fast)": "gpt2",
     "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
@@ -10,47 +10,59 @@ AVAILABLE_MODELS = {
 
 set_seed(42)
 
-# Cache loaded models to avoid reloading
-model_cache = {}
+# Cache models
+text_model_cache = {}
+image_generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v1-4")
 
-def load_model(model_name):
-    if model_name not in model_cache:
-        model_id = AVAILABLE_MODELS[model_name]
-        model_cache[model_name] = pipeline("text-generation", model=model_id)
-    return model_cache[model_name]
-
-# Chat history buffer
+# Session memory
 chat_memory = {}
 
-# Terminal response function
-def codette_terminal(user_input, model_name, session_id):
+def load_text_model(model_name):
+    if model_name not in text_model_cache:
+        text_model_cache[model_name] = pipeline("text-generation", model=AVAILABLE_MODELS[model_name])
+    return text_model_cache[model_name]
+
+def codette_terminal(prompt, model_name, generate_image, session_id):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
 
-    if user_input.lower() in ['exit', 'quit']:
+    if prompt.lower() in ["exit", "quit"]:
         chat_memory[session_id] = []
-        return "🧠 Codette signing off... Session reset."
+        return "🧠 Codette signing off... Session reset.", None
 
-    generator = load_model(model_name)
-    output = generator(user_input, max_length=100, num_return_sequences=1, do_sample=True)
-    response = output[0]['generated_text'].strip()
+    # Text generation
+    generator = load_text_model(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 > {user_input}")
+    chat_memory[session_id].append(f"🖋️ You > {prompt}")
     chat_memory[session_id].append(f"🧠 Codette > {response}")
-    return "\n".join(chat_memory[session_id][-10:])  # limit memory output
+    chat_log = "\n".join(chat_memory[session_id][-10:])
+
+    # Optional image generation
+    img = None
+    if generate_image:
+        img = image_generator(prompt)[0]['image']
+
+    return chat_log, img
 
-# UI
-with gr.Blocks(title="Codette Terminal - Multi-Model AI") as demo:
-    gr.Markdown("## 🧬 Codette Terminal (Multi-Model, Hugging Face Edition)")
-    gr.Markdown("Type something below. Use the dropdown to switch models. Type `'exit'` to reset session.")
+# Build Gradio UI
+with gr.Blocks(title="Codette Terminal: Text + Image AI") as demo:
+    gr.Markdown("## 🧬 Codette Terminal (Text + Image, Hugging Face Edition)")
+    gr.Markdown("Choose your model, enter a prompt. Enable image generation if desired. Type `'exit'` to reset.")
 
-    session_id = gr.Textbox(label="Session ID (hidden)", value="session_default", visible=False)
-    model_select = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Choose AI Model")
-    user_input = gr.Textbox(label="Your Prompt", placeholder="Ask me anything...")
-    output_box = gr.Textbox(label="Terminal Output", lines=20, interactive=False)
+    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)
+    user_input = gr.Textbox(label="Your Prompt", placeholder="e.g. a castle on Mars, explained by an AI philosopher", lines=1)
+    output_text = gr.Textbox(label="Terminal Output", lines=15, interactive=False)
+    output_image = gr.Image(label="AI-Generated Image")
 
-    user_input.submit(fn=codette_terminal, inputs=[user_input, model_select, session_id], outputs=output_box)
+    user_input.submit(
+        fn=codette_terminal,
+        inputs=[user_input, model_dropdown, generate_image_toggle, session_id],
+        outputs=[output_text, output_image]
+    )
 
-# For Hugging Face Spaces
+# Launch on Spaces
 if __name__ == "__main__":
     demo.launch()