Update app.py

app.py

@@ -2,17 +2,9 @@
 import os
 import base64
 import streamlit as st
-import pandas as pd
 import csv
 import time
 from dataclasses import dataclass
-from PIL import Image
-from datetime import datetime
-import pytz
-from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
-import av
-
-# Minimal initial imports to reduce startup delay
 
 st.set_page_config(page_title="SFT Tiny Titans 🚀", page_icon="🤖", layout="wide", initial_sidebar_state="expanded")
 
@@ -78,41 +70,43 @@ def get_download_link(file_path, mime_type="text/plain", label="Download"):
     return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
 
 def generate_filename(text_line):
+    from datetime import datetime
+    import pytz
     central = pytz.timezone('US/Central')
     timestamp = datetime.now(central).strftime("%Y%m%d_%I%M%S_%p")
     safe_text = ''.join(c if c.isalnum() else '_' for c in text_line[:50])
     return f"{timestamp}_{safe_text}.png"
 
 def get_gallery_files(file_types):
+    import glob
     return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
 
 # Video Transformer for WebRTC
-class VideoSnapshot(VideoTransformerBase):
+class VideoSnapshot:
     def __init__(self):
         self.snapshot = None
-    def transform(self, frame):
-        img = frame.to_ndarray(format="bgr24")
-        return img
+    def recv(self, frame):
+        from PIL import Image
+        img = frame.to_image()
+        self.snapshot = img
+        return frame
     def take_snapshot(self):
-        if self.snapshot is not None:
-            return Image.fromarray(self.snapshot)
+        return self.snapshot
 
 # Main App
-st.title("SFT Tiny Titans 🚀 (Lean & Mean!)")
+st.title("SFT Tiny Titans 🚀 (Fast & Furious!)")
 
 # Sidebar Galleries
 st.sidebar.header("Media Gallery 🎨")
-for gallery_type, file_types, emoji in [
-    ("Images 📸", ["png", "jpg", "jpeg"], "🖼️"),
-    ("Videos 🎥", ["mp4"], "🎬")
-]:
+for gallery_type, file_types, emoji in [("Images 📸", ["png", "jpg", "jpeg"], "🖼️"), ("Videos 🎥", ["mp4"], "🎬")]:
     st.sidebar.subheader(f"{gallery_type} {emoji}")
     files = get_gallery_files(file_types)
     if files:
-        cols = st.sidebar.columns(3)
-        for idx, file in enumerate(files[:6]):
-            with cols[idx % 3]:
+        cols = st.sidebar.columns(2)
+        for idx, file in enumerate(files[:4]):
+            with cols[idx % 2]:
                 if "Images" in gallery_type:
+                    from PIL import Image
                     st.image(Image.open(file), caption=file.split('/')[-1], use_column_width=True)
                 elif "Videos" in gallery_type:
                     st.video(file)
@@ -120,7 +114,7 @@ for gallery_type, file_types, emoji in [
 # Sidebar Model Management
 st.sidebar.subheader("Model Hub 🗂️")
 model_type = st.sidebar.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"])
-model_options = ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if "NLP" in model_type else ["stabilityai/stable-diffusion-2-1", "CompVis/stable-diffusion-v1-4"]
+model_options = ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if "NLP" in model_type else ["CompVis/stable-diffusion-v1-4"]
 selected_model = st.sidebar.selectbox("Select Model", ["None"] + model_options)
 if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
@@ -131,15 +125,10 @@ if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     st.session_state['model_loaded'] = True
 
 # Tabs
-tab1, tab2, tab3, tab4 = st.tabs([
-    "Build Titan 🌱",
-    "Fine-Tune Titans 🔧",
-    "Test Titans 🧪",
-    "Camera Snap 📷"
-])
+tab1, tab2, tab3, tab4 = st.tabs(["Build Titan 🌱", "Fine-Tune Titans 🔧", "Test Titans 🧪", "Camera Snap 📷"])
 
 with tab1:
-    st.header("Build Titan 🌱 (Start Small!)")
+    st.header("Build Titan 🌱 (Quick Start!)")
     model_type = st.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"], key="build_type")
     base_model = st.selectbox("Select Model", model_options, key="build_model")
     if st.button("Download Model ⬇️"):
@@ -149,10 +138,10 @@ with tab1:
             builder.load_model(base_model, config)
         st.session_state['builder'] = builder
         st.session_state['model_loaded'] = True
-        st.success("Titan ready! 🎉")
+        st.success("Titan up! 🎉")
 
 with tab2:
-    st.header("Fine-Tune Titans 🔧 (Sharpen Up!)")
+    st.header("Fine-Tune Titans 🔧 (Tune Fast!)")
     if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
         st.warning("Load a Titan first! ⚠️")
     else:
@@ -186,24 +175,26 @@ with tab2:
                 dataloader = DataLoader(dataset, batch_size=2)
                 optimizer = torch.optim.AdamW(st.session_state['builder'].model.parameters(), lr=2e-5)
                 st.session_state['builder'].model.train()
-                for _ in range(3):  # Simplified epochs
+                for _ in range(1):  # Minimal epochs
                     for batch in dataloader:
                         optimizer.zero_grad()
                         outputs = st.session_state['builder'].model(**{k: v.to(st.session_state['builder'].model.device) for k, v in batch.items()})
                         outputs.loss.backward()
                         optimizer.step()
-                st.success("NLP tuned! 🎉")
+                st.success("NLP sharpened! 🎉")
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
             st.subheader("CV Tune 🎨")
             uploaded_files = st.file_uploader("Upload Images", type=["png", "jpg"], accept_multiple_files=True, key="cv_upload")
             text_input = st.text_area("Text (one per image)", "Bat Neon\nIron Glow", key="cv_text")
             if uploaded_files and st.button("Tune CV 🔄"):
                 import torch
+                from PIL import Image
+                import numpy as np
                 images = [Image.open(f).convert("RGB") for f in uploaded_files]
                 texts = text_input.splitlines()[:len(images)]
                 optimizer = torch.optim.AdamW(st.session_state['builder'].pipeline.unet.parameters(), lr=1e-5)
                 st.session_state['builder'].pipeline.unet.train()
-                for _ in range(3):  # Simplified epochs
+                for _ in range(1):  # Minimal epochs
                     for img, text in zip(images, texts):
                         optimizer.zero_grad()
                         latents = st.session_state['builder'].pipeline.vae.encode(torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(st.session_state['builder'].pipeline.device)).latent_dist.sample()
@@ -218,16 +209,16 @@ with tab2:
                 for img, text in zip(images, texts):
                     filename = generate_filename(text)
                     img.save(filename)
-                st.success("CV tuned! 🎉")
+                st.success("CV polished! 🎉")
 
 with tab3:
-    st.header("Test Titans 🧪 (Showtime!)")
+    st.header("Test Titans 🧪 (Quick Check!)")
     if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
         st.warning("Load a Titan first! ⚠️")
     else:
         if isinstance(st.session_state['builder'], ModelBuilder):
             st.subheader("NLP Test 🧠")
-            prompt = st.text_area("Prompt", "What’s a superhero party?", key="nlp_test")
+            prompt = st.text_area("Prompt", "What’s a superhero?", key="nlp_test")
             if st.button("Test NLP ▶️"):
                 result = st.session_state['builder'].evaluate(prompt)
                 st.write(f"**Answer**: {result}")
@@ -240,12 +231,13 @@ with tab3:
                 st.image(img, caption="Generated Art")
 
 with tab4:
-    st.header("Camera Snap 📷 (Live Action!)")
-    ctx = webrtc_streamer(key="camera", video_transformer_factory=VideoSnapshot, rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})
-    if ctx.video_transformer:
+    st.header("Camera Snap 📷 (Instant Shots!)")
+    from streamlit_webrtc import webrtc_streamer
+    ctx = webrtc_streamer(key="camera", video_processor_factory=VideoSnapshot, rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})
+    if ctx.video_processor:
         snapshot_text = st.text_input("Snapshot Text", "Live Snap")
         if st.button("Snap It! 📸"):
-            snapshot = ctx.video_transformer.take_snapshot()
+            snapshot = ctx.video_processor.take_snapshot()
             if snapshot:
                 filename = generate_filename(snapshot_text)
                 snapshot.save(filename)
@@ -254,10 +246,11 @@ with tab4:
 
     # Demo Dataset
     st.subheader("Demo CV Dataset 🎨")
-    demo_texts = ["Bat Neon", "Iron Glow", "Thor Spark"]
+    demo_texts = ["Bat Neon", "Iron Glow"]
     demo_images = [generate_filename(t) for t in demo_texts]
     for img, text in zip(demo_images, demo_texts):
         if not os.path.exists(img):
+            from PIL import Image
             Image.new("RGB", (100, 100)).save(img)
     st.code("\n".join([f"{i+1}. {t} -> {img}" for i, (t, img) in enumerate(zip(demo_texts, demo_images))]), language="text")
     if st.button("Download Demo CSV 📝"):