Update app.py

app.py

@@ -42,6 +42,14 @@ st.set_page_config(
     }
 )
 
+# Initialize st.session_state
+if 'captured_images' not in st.session_state:
+    st.session_state['captured_images'] = []
+if 'builder' not in st.session_state:
+    st.session_state['builder'] = None
+if 'model_loaded' not in st.session_state:
+    st.session_state['model_loaded'] = False
+
 # Model Configuration Classes
 @dataclass
 class ModelConfig:
@@ -110,6 +118,7 @@ class ModelBuilder:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             if config:
                 self.config = config
+            self.model.to("cuda" if torch.cuda.is_available() else "cpu")
         st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
         return self
     def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
@@ -233,6 +242,15 @@ def get_model_files(model_type="causal_lm"):
 def get_gallery_files(file_types):
     return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
 
+def update_gallery():
+    media_files = get_gallery_files(["png"])
+    if media_files:
+        cols = st.sidebar.columns(2)
+        for idx, file in enumerate(media_files[:gallery_size * 2]):
+            with cols[idx % 2]:
+                st.image(Image.open(file), caption=file, use_container_width=True)
+                st.markdown(get_download_link(file, "image/png", "Download Image"), unsafe_allow_html=True)
+
 # Mock Search Tool for RAG
 def mock_search(query: str) -> str:
     if "superhero" in query.lower():
@@ -299,13 +317,7 @@ st.title("SFT Tiny Titans 🚀 (Small but Mighty!)")
 # Sidebar Galleries
 st.sidebar.header("Media Gallery 🎨")
 gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
-media_files = get_gallery_files(["png"])
-if media_files:
-    cols = st.sidebar.columns(2)
-    for idx, file in enumerate(media_files[:gallery_size * 2]):
-        with cols[idx % 2]:
-            st.image(Image.open(file), caption=file, use_column_width=True)
-            st.markdown(get_download_link(file, "image/png", "Download Image"), unsafe_allow_html=True)
+update_gallery()
 
 st.sidebar.subheader("Model Management 🗂️")
 model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"])
@@ -350,10 +362,8 @@ with tab2:
             filename = generate_filename(0)
             with open(filename, "wb") as f:
                 f.write(cam0_img.getvalue())
-            st.image(Image.open(filename), caption=filename, use_column_width=True)
+            st.image(Image.open(filename), caption=filename, use_container_width=True)
             logger.info(f"Saved snapshot from Camera 0: {filename}")
-            if 'captured_images' not in st.session_state:
-                st.session_state['captured_images'] = []
             st.session_state['captured_images'].append(filename)
             update_gallery()
         if st.button(f"Capture {slice_count} Frames - Cam 0 📸"):
@@ -370,7 +380,7 @@ with tab2:
             st.session_state['captured_images'].extend(st.session_state['cam0_frames'])
             update_gallery()
             for frame in st.session_state['cam0_frames']:
-                st.image(Image.open(frame), caption=frame, use_column_width=True)
+                st.image(Image.open(frame), caption=frame, use_container_width=True)
     with cols[1]:
         st.subheader("Camera 1")
         cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
@@ -378,10 +388,8 @@ with tab2:
             filename = generate_filename(1)
             with open(filename, "wb") as f:
                 f.write(cam1_img.getvalue())
-            st.image(Image.open(filename), caption=filename, use_column_width=True)
+            st.image(Image.open(filename), caption=filename, use_container_width=True)
             logger.info(f"Saved snapshot from Camera 1: {filename}")
-            if 'captured_images' not in st.session_state:
-                st.session_state['captured_images'] = []
             st.session_state['captured_images'].append(filename)
             update_gallery()
         if st.button(f"Capture {slice_count} Frames - Cam 1 📸"):
@@ -398,7 +406,7 @@ with tab2:
             st.session_state['captured_images'].extend(st.session_state['cam1_frames'])
             update_gallery()
             for frame in st.session_state['cam1_frames']:
-                st.image(Image.open(frame), caption=frame, use_column_width=True)
+                st.image(Image.open(frame), caption=frame, use_container_width=True)
 
 with tab3:
     st.header("Fine-Tune Titan 🔧")
@@ -485,4 +493,7 @@ st.sidebar.subheader("Action Logs 📜")
 log_container = st.sidebar.empty()
 with log_container:
     for record in log_records:
-        st.write(f"{record.asctime} - {record.levelname} - {record.message}")
+        st.write(f"{record.asctime} - {record.levelname} - {record.message}")
+
+# Initial Gallery Update
+update_gallery()