Update app.py

app.py

@@ -20,6 +20,7 @@ chat_memory = {}
 try:
     image_generator = DiffusionPipeline.from_pretrained(
         "runwayml/stable-diffusion-v1-5",
+        safety_checker=None,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     )
     image_generator.to(device)
@@ -33,6 +34,7 @@ except Exception as e:
 try:
     video_pipeline = DiffusionPipeline.from_pretrained(
         "damo-vilab/text-to-video-ms-1.7b",
+        safety_checker=None,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     )
     video_pipeline.to(device)
@@ -43,7 +45,7 @@ except Exception as e:
     video_enabled = False
 
 # ---------- Streamed Response Generator ----------
-def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
+def codette_terminal(prompt, model_name, generate_image, generate_video, session_id, batch_size, video_steps, fps):
     if session_id not in chat_memory:
         chat_memory[session_id] = []
 
@@ -89,10 +91,11 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
     chat_memory[session_id].append(f"🧠 Codette > {output}")
 
     # Image Generation
-    img = None
+    imgs = None
     if generate_image and image_enabled:
         try:
-            img = image_generator(prompt).images[0]
+            result = image_generator(prompt, num_images_per_prompt=batch_size)
+            imgs = result.images
         except Exception as e:
             response_so_far += f"\n[Image error]: {e}"
 
@@ -100,32 +103,36 @@ def codette_terminal(prompt, model_name, generate_image, generate_video, session
     vid = None
     if generate_video and video_enabled:
         try:
-            video_frames = video_pipeline(prompt, num_inference_steps=50).frames
+            result = video_pipeline(prompt, num_inference_steps=video_steps)
+            frames = result.frames
             temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-            imageio.mimsave(temp_video_path, video_frames, fps=8)
+            imageio.mimsave(temp_video_path, frames, fps=fps)
             vid = temp_video_path
         except Exception as e:
             response_so_far += f"\n[Video error]: {e}"
 
-    yield "\n".join(chat_memory[session_id][-10:]), img, vid
+    yield "\n".join(chat_memory[session_id][-10:]), imgs, vid
 
 # ---------- Gradio UI ----------
 with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
-    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video)")
-    gr.Markdown("Type a prompt and select your model. Enable image or video generation if you like. Type `'exit'` to reset.")
+    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video + Batch + NSFW OK)")
+    gr.Markdown("Type a prompt, select your model, and configure generation options. Type `'exit'` to reset.")
 
     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 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)
+    model_dropdown = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Language Model")
+    generate_image_toggle = gr.Checkbox(label="Generate Image(s)?", value=False, interactive=image_enabled)
+    generate_video_toggle = gr.Checkbox(label="Generate Video?", value=False, interactive=video_enabled)
+    batch_size_slider = gr.Slider(label="Number of Images", minimum=1, maximum=4, step=1, value=1)
+    video_steps_slider = gr.Slider(label="Video Inference Steps", minimum=10, maximum=100, step=10, value=50)
+    fps_slider = gr.Slider(label="Video FPS", minimum=4, maximum=24, step=2, value=8)
     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_image = gr.Gallery(label="Generated Image(s)").style(grid=2)
     output_video = gr.Video(label="Generated Video")
 
     user_input.submit(
         codette_terminal,
-        inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
+        inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id, batch_size_slider, video_steps_slider, fps_slider],
         outputs=[output_text, output_image, output_video],
         concurrency_limit=1,
         queue=True,