Raiff1982 commited on
Commit
d90bae9
·
verified ·
1 Parent(s): c160902

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -52
app.py CHANGED
@@ -1,36 +1,18 @@
1
- import gradio as gr
2
- from transformers import pipeline, set_seed
3
 
4
- # Define available models
5
- AVAILABLE_MODELS = {
6
- "GPT-2 (small, fast)": "gpt2",
7
- "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
8
- "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
9
- }
10
 
11
- set_seed(42)
12
-
13
- # Cache models
14
- text_model_cache = {}
15
- image_generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v1-4")
16
-
17
- # Session memory
18
- chat_memory = {}
19
-
20
- def load_text_model(model_name):
21
- if model_name not in text_model_cache:
22
- text_model_cache[model_name] = pipeline("text-generation", model=AVAILABLE_MODELS[model_name])
23
- return text_model_cache[model_name]
24
-
25
- def codette_terminal(prompt, model_name, generate_image, session_id):
26
  if session_id not in chat_memory:
27
  chat_memory[session_id] = []
28
 
29
  if prompt.lower() in ["exit", "quit"]:
30
  chat_memory[session_id] = []
31
- return "🧠 Codette signing off... Session reset.", None
32
 
33
- # Text generation
34
  generator = load_text_model(model_name)
35
  response = generator(prompt, max_length=100, num_return_sequences=1, do_sample=True)[0]['generated_text'].strip()
36
 
@@ -38,31 +20,35 @@ def codette_terminal(prompt, model_name, generate_image, session_id):
38
  chat_memory[session_id].append(f"🧠 Codette > {response}")
39
  chat_log = "\n".join(chat_memory[session_id][-10:])
40
 
41
- # Optional image generation
42
- img = None
 
43
  if generate_image:
44
- img = image_generator(prompt)[0]['image']
45
-
46
- return chat_log, img
47
-
48
- # Build Gradio UI
49
- with gr.Blocks(title="Codette Terminal: Text + Image AI") as demo:
50
- gr.Markdown("## 🧬 Codette Terminal (Text + Image, Hugging Face Edition)")
51
- gr.Markdown("Choose your model, enter a prompt. Enable image generation if desired. Type `'exit'` to reset.")
52
-
53
- session_id = gr.Textbox(value="session_default", visible=False)
54
- model_dropdown = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Choose Language Model")
55
- generate_image_toggle = gr.Checkbox(label="Also generate image?", value=False)
56
- user_input = gr.Textbox(label="Your Prompt", placeholder="e.g. a castle on Mars, explained by an AI philosopher", lines=1)
57
- output_text = gr.Textbox(label="Terminal Output", lines=15, interactive=False)
58
- output_image = gr.Image(label="AI-Generated Image")
59
-
60
- user_input.submit(
61
- fn=codette_terminal,
62
- inputs=[user_input, model_dropdown, generate_image_toggle, session_id],
63
- outputs=[output_text, output_image]
64
- )
65
-
66
- # Launch on Spaces
67
- if __name__ == "__main__":
68
- demo.launch()
 
 
 
 
1
+ from diffusers import DiffusionPipeline
2
+ import tempfile
3
 
4
+ # Load video generation model (you can switch to others)
5
+ video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16)
6
+ video_pipeline = video_pipeline.to("cuda")
 
 
 
7
 
8
+ def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
  if session_id not in chat_memory:
10
  chat_memory[session_id] = []
11
 
12
  if prompt.lower() in ["exit", "quit"]:
13
  chat_memory[session_id] = []
14
+ return "🧠 Codette signing off... Session reset.", None, None
15
 
 
16
  generator = load_text_model(model_name)
17
  response = generator(prompt, max_length=100, num_return_sequences=1, do_sample=True)[0]['generated_text'].strip()
18
 
 
20
  chat_memory[session_id].append(f"🧠 Codette > {response}")
21
  chat_log = "\n".join(chat_memory[session_id][-10:])
22
 
23
+ image = None
24
+ video = None
25
+
26
  if generate_image:
27
+ image = image_generator(prompt)[0]['image']
28
+
29
+ if generate_video:
30
+ # Generate video from prompt
31
+ video_frames = video_pipeline(prompt, num_inference_steps=50).frames
32
+ # Save frames as video
33
+ temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
34
+ import imageio
35
+ imageio.mimsave(temp_video_path, video_frames, fps=8)
36
+ video = temp_video_path
37
+
38
+ return chat_log, image, video
39
+
40
+
41
+
42
+ generate_video_toggle = gr.Checkbox(label="Also generate video?", value=False)
43
+ output_video = gr.Video(label="AI-Generated Video")
44
+
45
+ user_input.submit(
46
+ fn=codette_terminal,
47
+ inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
48
+ outputs=[output_text, output_image, output_video]
49
+ )
50
+
51
+
52
+
53
+
54
+