Update app.py

app.py

@@ -175,17 +175,17 @@ class DiffusionBuilder:
         self.config = None
         self.pipeline = None
         self.model_type = None
-    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None, model_type: str = "StableDiffusion"):
-        with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
+    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None, model_type: str = "StableDiffusion", download: bool = True):
+        with st.spinner(f"{'Downloading' if download else 'Loading'} {model_path}... ⏳"):
             if model_type == "StableDiffusion":
-                self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
+                self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32, use_safetensors=True, local_files_only=not download).to("cpu")
             elif model_type == "DDPM":
-                self.pipeline = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
+                self.pipeline = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch.float32, use_safetensors=True, local_files_only=not download).to("cpu")
                 self.pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipeline.scheduler.config)
             if config:
                 self.config = config
             self.model_type = model_type
-        st.success(f"Diffusion model loaded! 🎨")
+        st.success(f"Diffusion model {'downloaded' if download else 'loaded'}! 🎨")
         return self
     def fine_tune_sft(self, images, texts, epochs=3):
         dataset = DiffusionDataset(images, texts)
@@ -339,11 +339,35 @@ if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     st.session_state['model_loaded'] = True
     st.rerun()
 
-# Tabs (Reordered: Camera Snap first)
-tab1, tab2, tab3, tab4, tab5 = st.tabs(["Camera Snap 📷", "Fine-Tune Titan 🔧", "Build Titan 🌱", "Test Titan 🧪", "Agentic RAG Party 🌐"])
+# Tabs
+tab1, tab2, tab3, tab4 = st.tabs(["Build Titan & Camera Snap 🌱📷", "Fine-Tune Titan 🔧", "Test Titan 🧪", "Agentic RAG Party 🌐"])
 
 with tab1:
-    st.header("Camera Snap 📷 (Dual Capture!)")
+    st.header("Build Titan & Camera Snap 🌱📷")
+    st.subheader("Build Titan 🌱")
+    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
+    base_model_options = {
+        "Causal LM": ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"],
+        "Diffusion": [
+            "OFA-Sys/small-stable-diffusion-v0 (LDM/Conditional, ~300 MB)", 
+            "google/ddpm-ema-celebahq-256 (DDPM/SDE/Autoregressive Proxy, ~280 MB)"
+        ]
+    }
+    base_model = st.selectbox("Select Tiny Model", base_model_options[model_type])
+    action = st.radio("Action", ["Use Model", "Download Model"], index=0 if "Causal LM" in model_type else 1)
+    model_name = st.text_input("Model Name (for Download)", f"tiny-titan-{int(time.time())}") if action == "Download Model" else None
+    if st.button(f"{action} ⬇️"):
+        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name or base_model.split(" ")[0], base_model=base_model.split(" ")[0], size="small")
+        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
+        model_type_for_diffusion = "StableDiffusion" if "small-stable-diffusion" in base_model else "DDPM"
+        builder.load_model(base_model.split(" ")[0], config, model_type_for_diffusion, download=action == "Download Model")
+        if action == "Download Model":
+            builder.save_model(config.model_path)
+        st.session_state['builder'] = builder
+        st.session_state['model_loaded'] = True
+        st.rerun()
+
+    st.subheader("Camera Snap 📷")
     slice_count = st.number_input("Image Slice Count", min_value=1, max_value=20, value=10)
     video_length = st.number_input("Video Length (seconds)", min_value=1, max_value=30, value=10)
     cols = st.columns(2)
@@ -352,24 +376,26 @@ with tab1:
         cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
         if cam0_img:
             filename = generate_filename(0)
-            with open(filename, "wb") as f:
-                f.write(cam0_img.getvalue())
-            st.image(Image.open(filename), caption=filename, use_container_width=True)
-            logger.info(f"Saved snapshot from Camera 0: {filename}")
-            st.session_state['captured_images'].append(filename)
-            update_gallery()
+            if filename not in st.session_state['captured_images']:
+                with open(filename, "wb") as f:
+                    f.write(cam0_img.getvalue())
+                st.image(Image.open(filename), caption=filename, use_container_width=True)
+                logger.info(f"Saved snapshot from Camera 0: {filename}")
+                st.session_state['captured_images'].append(filename)
+                update_gallery()
         if st.button(f"Capture {slice_count} Frames - Cam 0 📸"):
             st.session_state['cam0_frames'] = []
             for i in range(slice_count):
                 img = st.camera_input(f"Frame {i} - Cam 0", key=f"cam0_frame_{i}_{time.time()}")
                 if img:
                     filename = generate_filename(f"0_{i}")
-                    with open(filename, "wb") as f:
-                        f.write(img.getvalue())
-                    st.session_state['cam0_frames'].append(filename)
-                    logger.info(f"Saved frame {i} from Camera 0: {filename}")
-                    time.sleep(1.0 / slice_count)
-            st.session_state['captured_images'].extend(st.session_state['cam0_frames'])
+                    if filename not in st.session_state['captured_images']:
+                        with open(filename, "wb") as f:
+                            f.write(img.getvalue())
+                        st.session_state['cam0_frames'].append(filename)
+                        logger.info(f"Saved frame {i} from Camera 0: {filename}")
+                        time.sleep(1.0 / slice_count)
+            st.session_state['captured_images'].extend([f for f in st.session_state['cam0_frames'] if f not in st.session_state['captured_images']])
             update_gallery()
             for frame in st.session_state['cam0_frames']:
                 st.image(Image.open(frame), caption=frame, use_container_width=True)
@@ -378,24 +404,26 @@ with tab1:
         cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
         if cam1_img:
             filename = generate_filename(1)
-            with open(filename, "wb") as f:
-                f.write(cam1_img.getvalue())
-            st.image(Image.open(filename), caption=filename, use_container_width=True)
-            logger.info(f"Saved snapshot from Camera 1: {filename}")
-            st.session_state['captured_images'].append(filename)
-            update_gallery()
+            if filename not in st.session_state['captured_images']:
+                with open(filename, "wb") as f:
+                    f.write(cam1_img.getvalue())
+                st.image(Image.open(filename), caption=filename, use_container_width=True)
+                logger.info(f"Saved snapshot from Camera 1: {filename}")
+                st.session_state['captured_images'].append(filename)
+                update_gallery()
         if st.button(f"Capture {slice_count} Frames - Cam 1 📸"):
             st.session_state['cam1_frames'] = []
             for i in range(slice_count):
                 img = st.camera_input(f"Frame {i} - Cam 1", key=f"cam1_frame_{i}_{time.time()}")
                 if img:
                     filename = generate_filename(f"1_{i}")
-                    with open(filename, "wb") as f:
-                        f.write(img.getvalue())
-                    st.session_state['cam1_frames'].append(filename)
-                    logger.info(f"Saved frame {i} from Camera 1: {filename}")
-                    time.sleep(1.0 / slice_count)
-            st.session_state['captured_images'].extend(st.session_state['cam1_frames'])
+                    if filename not in st.session_state['captured_images']:
+                        with open(filename, "wb") as f:
+                            f.write(img.getvalue())
+                        st.session_state['cam1_frames'].append(filename)
+                        logger.info(f"Saved frame {i} from Camera 1: {filename}")
+                        time.sleep(1.0 / slice_count)
+            st.session_state['captured_images'].extend([f for f in st.session_state['cam1_frames'] if f not in st.session_state['captured_images']])
             update_gallery()
             for frame in st.session_state['cam1_frames']:
                 st.image(Image.open(frame), caption=frame, use_container_width=True)
@@ -444,28 +472,6 @@ with tab2:
                 st.markdown(get_download_link(csv_path, "text/csv", "Download SFT Dataset CSV"), unsafe_allow_html=True)
 
 with tab3:
-    st.header("Build Titan 🌱")
-    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
-    base_model_options = {
-        "Causal LM": ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"],
-        "Diffusion": [
-            "OFA-Sys/small-stable-diffusion-v0 (LDM/Conditional)", 
-            "google/ddpm-ema-celebahq-256 (DDPM/SDE/Autoregressive Proxy)"
-        ]
-    }
-    base_model = st.selectbox("Select Tiny Model", base_model_options[model_type])
-    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
-    if st.button("Download Model ⬇️"):
-        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model.split(" ")[0], size="small")
-        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
-        model_type_for_diffusion = "StableDiffusion" if "small-stable-diffusion" in base_model else "DDPM"
-        builder.load_model(base_model.split(" ")[0], config, model_type_for_diffusion)
-        builder.save_model(config.model_path)
-        st.session_state['builder'] = builder
-        st.session_state['model_loaded'] = True
-        st.rerun()
-
-with tab4:
     st.header("Test Titan 🧪")
     if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
         st.warning("Please build or load a Titan first! ⚠️")
@@ -487,7 +493,7 @@ with tab4:
                         image = st.session_state['builder'].generate(prompt)
                     st.image(image, caption=f"Generated from {selected_pipeline}")
 
-with tab5:
+with tab4:
     st.header("Agentic RAG Party 🌐")
     if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
         st.warning("Please build or load a Titan first! ⚠️")