Update app.py

app.py

@@ -81,7 +81,7 @@ 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
+# Video Processor for WebRTC
 class VideoSnapshot:
     def __init__(self):
         self.snapshot = None
@@ -94,7 +94,7 @@ class VideoSnapshot:
         return self.snapshot
 
 # Main App
-st.title("SFT Tiny Titans 🚀 (Fast & Furious!)")
+st.title("SFT Tiny Titans 🚀 (Fast & Fixed!)")
 
 # Sidebar Galleries
 st.sidebar.header("Media Gallery 🎨")
@@ -107,15 +107,15 @@ for gallery_type, file_types, emoji in [("Images 📸", ["png", "jpg", "jpeg"],
             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)
+                    st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)
                 elif "Videos" in gallery_type:
                     st.video(file)
 
 # 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 ["CompVis/stable-diffusion-v1-4"]
-selected_model = st.sidebar.selectbox("Select Model", ["None"] + model_options)
+model_options = {"NLP (Causal LM)": "HuggingFaceTB/SmolLM-135M", "CV (Diffusion)": "CompVis/stable-diffusion-v1-4"}
+selected_model = st.sidebar.selectbox("Select Model", ["None", model_options[model_type]])
 if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
     config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=selected_model)
@@ -130,7 +130,7 @@ tab1, tab2, tab3, tab4 = st.tabs(["Build Titan 🌱", "Fine-Tune Titans 🔧", "
 with tab1:
     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")
+    base_model = st.selectbox("Select Model", [model_options[model_type]], key="build_model")
     if st.button("Download Model ⬇️"):
         config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=base_model)
         builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
@@ -175,7 +175,7 @@ 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(1):  # Minimal epochs
+                for _ in range(1):
                     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()})
@@ -194,7 +194,7 @@ with tab2:
                 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(1):  # Minimal epochs
+                for _ in range(1):
                     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()
@@ -233,7 +233,11 @@ with tab3:
 with tab4:
     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"]}]})
+    ctx = webrtc_streamer(
+        key="camera",
+        video_processor_factory=VideoSnapshot,
+        frontend_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! 📸"):
@@ -241,7 +245,7 @@ with tab4:
             if snapshot:
                 filename = generate_filename(snapshot_text)
                 snapshot.save(filename)
-                st.image(snapshot, caption=filename)
+                st.image(snapshot, caption=filename, use_container_width=True)
                 st.success("Snapped! 🎉")
 
     # Demo Dataset