Update app.py

app.py

@@ -54,7 +54,7 @@ st.set_page_config(
 
 # Initialize st.session_state
 if 'captured_files' not in st.session_state:
-    st.session_state['captured_files'] = {'cam0': None, 'cam1': None}  # Dictionary to store one file per camera
+    st.session_state['captured_files'] = {'cam0': None, 'cam1': None}  # One file per camera
 if 'builder' not in st.session_state:
     st.session_state['builder'] = None
 if 'model_loaded' not in st.session_state:
@@ -62,7 +62,7 @@ if 'model_loaded' not in st.session_state:
 if 'processing' not in st.session_state:
     st.session_state['processing'] = {}
 if 'history' not in st.session_state:
-    st.session_state['history'] = []
+    st.session_state['history'] = {'cam0': None, 'cam1': None}  # One history entry per camera
 
 # Model Configuration Classes
 @dataclass
@@ -506,20 +506,15 @@ st.title("AI Vision & SFT Titans 🚀")
 
 # Sidebar
 st.sidebar.header("Captured Files 📜")
-gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
+gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2)  # Default to 2 for two cameras
 def update_gallery():
-    media_files = list(st.session_state['captured_files'].values())
-    if media_files:
+    media_files = [st.session_state['captured_files']['cam0'], st.session_state['captured_files']['cam1']]
+    valid_files = [f for f in media_files if f and os.path.exists(f)]  # Only valid files
+    if valid_files:
         cols = st.sidebar.columns(2)
-        for idx, file in enumerate(media_files[:gallery_size * 2]):
-            if file and os.path.exists(file):  # Check if file exists
-                with cols[idx % 2]:
-                    if file.endswith(".png"):
-                        st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
-                    elif file.endswith(".txt"):
-                        with open(file, "r") as f:
-                            content = f.read()
-                            st.text(content[:50] + "..." if len(content) > 50 else content)
+        for idx, file in enumerate(valid_files[:2]):  # Limit to 2 images
+            with cols[idx]:
+                st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
 update_gallery()
 
 st.sidebar.subheader("Model Management 🗂️")
@@ -543,7 +538,8 @@ with log_container:
 st.sidebar.subheader("History 📜")
 history_container = st.sidebar.empty()
 with history_container:
-    for entry in st.session_state['history'][-5:]:
+    valid_history = [st.session_state['history']['cam0'], st.session_state['history']['cam1']]
+    for entry in [e for e in valid_history if e][:2]:  # Show only latest two
         st.write(entry)
 
 # Tabs
@@ -563,9 +559,9 @@ with tab1:
             with open(filename, "wb") as f:
                 f.write(cam0_img.getvalue())
             st.session_state['captured_files']['cam0'] = filename
+            st.session_state['history']['cam0'] = f"Snapshot from Cam 0: {filename}"
             st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
             logger.info(f"Saved snapshot from Camera 0: {filename}")
-            st.session_state['history'].append(f"Snapshot from Cam 0: {filename}")
             update_gallery()
     with cols[1]:
         cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
@@ -574,9 +570,9 @@ with tab1:
             with open(filename, "wb") as f:
                 f.write(cam1_img.getvalue())
             st.session_state['captured_files']['cam1'] = filename
+            st.session_state['history']['cam1'] = f"Snapshot from Cam 1: {filename}"
             st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
             logger.info(f"Saved snapshot from Camera 1: {filename}")
-            st.session_state['history'].append(f"Snapshot from Cam 1: {filename}")
             update_gallery()
 
 with tab2:
@@ -590,7 +586,7 @@ with tab2:
                 pdf_path = generate_filename("downloaded", "pdf")
                 if download_pdf(url, pdf_path):
                     logger.info(f"Downloaded PDF from {url} to {pdf_path}")
-                    st.session_state['history'].append(f"Downloaded PDF: {pdf_path}")
+                    st.session_state['history']['pdf'] = f"Downloaded PDF: {pdf_path}"
                     snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode.lower().replace(" ", "")))
                     for snapshot in snapshots:
                         st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
@@ -612,7 +608,7 @@ with tab3:
         builder.save_model(config.model_path)
         st.session_state['builder'] = builder
         st.session_state['model_loaded'] = True
-        st.session_state['history'].append(f"Built {model_type} model: {model_name}")
+        st.session_state['history']['build'] = f"Built {model_type} model: {model_name}"
         st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
         st.rerun()
 
@@ -647,7 +643,7 @@ with tab4:
                 st.session_state['builder'].save_model(new_config.model_path)
                 zip_path = f"{new_config.model_path}.zip"
                 zip_directory(new_config.model_path, zip_path)
-                st.session_state['history'].append(f"Fine-tuned Causal LM: {new_model_name}")
+                st.session_state['history']['sft'] = f"Fine-tuned Causal LM: {new_model_name}"
                 st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
                 st.rerun()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
@@ -665,7 +661,7 @@ with tab4:
                     st.session_state['builder'].save_model(new_config.model_path)
                     zip_path = f"{new_config.model_path}.zip"
                     zip_directory(new_config.model_path, zip_path)
-                    st.session_state['history'].append(f"Fine-tuned Diffusion: {new_model_name}")
+                    st.session_state['history']['sft'] = f"Fine-tuned Diffusion: {new_model_name}"
                     st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Diffusion Model"), unsafe_allow_html=True)
                 csv_path = f"sft_dataset_{int(time.time())}.csv"
                 with open(csv_path, "w", newline="") as f:
@@ -697,7 +693,7 @@ with tab5:
             if st.button("Run Test ▶️"):
                 status_container = st.empty()
                 result = st.session_state['builder'].evaluate(test_prompt, status_container)
-                st.session_state['history'].append(f"Causal LM Test: {test_prompt} -> {result}")
+                st.session_state['history']['test'] = f"Causal LM Test: {test_prompt} -> {result}"
                 st.write(f"**Generated Response**: {result}")
                 status_container.empty()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
@@ -707,7 +703,7 @@ with tab5:
                 output_file = generate_filename("diffusion_test", "png")
                 image.save(output_file)
                 st.session_state['captured_files']['diffusion_test'] = output_file
-                st.session_state['history'].append(f"Diffusion Test: {test_prompt} -> {output_file}")
+                st.session_state['history']['test'] = f"Diffusion Test: {test_prompt} -> {output_file}"
                 st.image(image, caption="Generated Image")
                 update_gallery()
 
@@ -721,14 +717,14 @@ with tab6:
                 agent = PartyPlannerAgent(st.session_state['builder'].model, st.session_state['builder'].tokenizer)
                 task = "Plan a luxury superhero-themed party at Wayne Manor."
                 plan_df = agent.plan_party(task)
-                st.session_state['history'].append(f"NLP RAG Demo: Planned party at Wayne Manor")
+                st.session_state['history']['rag'] = f"NLP RAG Demo: Planned party at Wayne Manor"
                 st.dataframe(plan_df)
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
             if st.button("Run CV RAG Demo 🎉"):
                 agent = CVPartyPlannerAgent(st.session_state['builder'].pipeline)
                 task = "Generate images for a luxury superhero-themed party."
                 plan_df = agent.plan_party(task)
-                st.session_state['history'].append(f"CV RAG Demo: Generated party images")
+                st.session_state['history']['rag'] = f"CV RAG Demo: Generated party images"
                 st.dataframe(plan_df)
                 for _, row in plan_df.iterrows():
                     image = agent.generate(row["Image Idea"])
@@ -750,7 +746,7 @@ with tab7:
                 output_file = generate_filename("ocr_output", "txt")
                 st.session_state['processing']['ocr'] = True
                 result = asyncio.run(process_ocr(image, output_file))
-                st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")
+                st.session_state['history']['ocr'] = f"OCR Test: {selected_file} -> {output_file}"
                 st.text_area("OCR Result", result, height=200, key="ocr_result")
                 st.success(f"OCR output saved to {output_file}")
                 st.session_state['processing']['ocr'] = False
@@ -770,7 +766,7 @@ with tab8:
                 output_file = generate_filename("gen_output", "png")
                 st.session_state['processing']['gen'] = True
                 result = asyncio.run(process_image_gen(prompt, output_file))
-                st.session_state['history'].append(f"Image Gen Test: {prompt} -> {output_file}")
+                st.session_state['history']['gen'] = f"Image Gen Test: {prompt} -> {output_file}"
                 st.image(result, caption="Generated Image", use_container_width=True)
                 st.success(f"Image saved to {output_file}")
                 st.session_state['processing']['gen'] = False
@@ -805,7 +801,7 @@ with tab9:
                 result = builder.generate("A superhero scene inspired by captured images")
                 result.save(output_file)
                 st.session_state['captured_files']['diffusion'] = output_file
-            st.session_state['history'].append(f"Custom Diffusion: {model_choice} -> {output_file}")
+            st.session_state['history']['diffusion'] = f"Custom Diffusion: {model_choice} -> {output_file}"
             st.image(result, caption=f"{model_choice} Masterpiece", use_container_width=True)
             st.success(f"Image saved to {output_file}")
             st.session_state['processing']['diffusion'] = False