Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import pipeline, set_seed | |
| from diffusers import DiffusionPipeline | |
| import tempfile | |
| 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] = [] | |
| if prompt.lower() in ["exit", "quit"]: | |
| chat_memory[session_id] = [] | |
| return "🧠 Codette signing off... Session reset.", None, None | |
| # 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 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() | |