Update app.py

app.py

@@ -53,10 +53,8 @@ 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}  # One file per camera
 if 'history' not in st.session_state:
-    st.session_state['history'] = {'cam0': None, 'cam1': None}  # One history entry per camera
+    st.session_state['history'] = []  # Flat list for history
 if 'builder' not in st.session_state:
     st.session_state['builder'] = None
 if 'model_loaded' not in st.session_state:
@@ -329,21 +327,9 @@ def get_model_files(model_type="causal_lm"):
     path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
     return [d for d in glob.glob(path) if os.path.isdir(d)]
 
-def get_gallery_files(file_types):
+def get_gallery_files(file_types=["png"]):
     return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
 
-def download_pdf(url, output_path):
-    try:
-        response = requests.get(url, stream=True, timeout=10)
-        if response.status_code == 200:
-            with open(output_path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
-            return True
-    except requests.RequestException as e:
-        logger.error(f"Failed to download {url}: {e}")
-    return False
-
 # Mock Search Tool for RAG
 def mock_search(query: str) -> str:
     if "superhero" in query.lower():
@@ -445,9 +431,6 @@ async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
                 output_files.append(output_file)
         elapsed = int(time.time() - start_time)
         status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
-        for file in output_files:
-            if file not in st.session_state['captured_files'].values():
-                st.session_state['captured_files'][f"pdf_{len(output_files)}"] = file
         update_gallery()
         return output_files
     except Exception as e:
@@ -465,8 +448,6 @@ async def process_ocr(image, output_file):
     status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
     async with aiofiles.open(output_file, "w") as f:
         await f.write(result)
-    if output_file not in st.session_state['captured_files'].values():
-        st.session_state['captured_files']['ocr'] = output_file
     update_gallery()
     return result
 
@@ -479,8 +460,6 @@ async def process_image_gen(prompt, output_file):
     elapsed = int(time.time() - start_time)
     status.text(f"Image Gen completed in {elapsed}s!")
     gen_image.save(output_file)
-    if output_file not in st.session_state['captured_files'].values():
-        st.session_state['captured_files']['gen'] = output_file
     update_gallery()
     return gen_image
 
@@ -496,8 +475,6 @@ async def process_custom_diffusion(images, output_file, model_name):
     elapsed = int(time.time() - start_time)
     status.text(f"{model_name} completed in {elapsed}s!")
     upscaled_image.save(output_file)
-    if output_file not in st.session_state['captured_files'].values():
-        st.session_state['captured_files']['diffusion'] = output_file
     update_gallery()
     return upscaled_image
 
@@ -506,18 +483,14 @@ st.title("AI Vision & SFT Titans 🚀")
 
 # Sidebar
 st.sidebar.header("Captured Files 📜")
-gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2)  # Default to 2 for two cameras
+gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2)  # Default to 2
 def update_gallery():
-    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:
+    media_files = get_gallery_files(["png"])
+    if media_files:
         cols = st.sidebar.columns(2)
-        if st.session_state['captured_files']['cam0'] in valid_files:
-            with cols[0]:
-                st.image(Image.open(st.session_state['captured_files']['cam0']), caption="Camera 0", use_container_width=True)
-        if st.session_state['captured_files']['cam1'] in valid_files:
-            with cols[1]:
-                st.image(Image.open(st.session_state['captured_files']['cam1']), caption="Camera 1", use_container_width=True)
+        for idx, file in enumerate(media_files[:gallery_size * 2]):  # Limit by gallery size
+            with cols[idx % 2]:
+                st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
 update_gallery()
 
 st.sidebar.subheader("Model Management 🗂️")
@@ -541,8 +514,7 @@ with log_container:
 st.sidebar.subheader("History 📜")
 history_container = st.sidebar.empty()
 with history_container:
-    valid_history = [st.session_state['history']['cam0'], st.session_state['history']['cam1']]
-    for entry in [e for e in valid_history if e]:  # Show only non-None entries
+    for entry in st.session_state['history'][-gallery_size * 2:]:  # Limit by gallery size
         st.write(entry)
 
 # Tabs
@@ -561,8 +533,9 @@ with tab1:
             filename = generate_filename("cam0")
             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}"
+            entry = f"Snapshot from Cam 0: {filename}"
+            if entry not in st.session_state['history']:
+                st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 0:")] + [entry]
             st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
             logger.info(f"Saved snapshot from Camera 0: {filename}")
             update_gallery()
@@ -572,8 +545,9 @@ with tab1:
             filename = generate_filename("cam1")
             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}"
+            entry = f"Snapshot from Cam 1: {filename}"
+            if entry not in st.session_state['history']:
+                st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 1:")] + [entry]
             st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
             logger.info(f"Saved snapshot from Camera 1: {filename}")
             update_gallery()
@@ -589,7 +563,9 @@ 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']['pdf'] = f"Downloaded PDF: {pdf_path}"
+                    entry = f"Downloaded PDF: {pdf_path}"
+                    if entry not in st.session_state['history']:
+                        st.session_state['history'].append(entry)
                     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)
@@ -611,7 +587,9 @@ with tab3:
         builder.save_model(config.model_path)
         st.session_state['builder'] = builder
         st.session_state['model_loaded'] = True
-        st.session_state['history']['build'] = f"Built {model_type} model: {model_name}"
+        entry = f"Built {model_type} model: {model_name}"
+        if entry not in st.session_state['history']:
+            st.session_state['history'].append(entry)
         st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
         st.rerun()
 
@@ -646,13 +624,15 @@ 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']['sft'] = f"Fine-tuned Causal LM: {new_model_name}"
+                entry = f"Fine-tuned Causal LM: {new_model_name}"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 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):
-            captured_files = list(st.session_state['captured_files'].values())
+            captured_files = get_gallery_files(["png"])
             if len(captured_files) >= 2:
-                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files if img]
+                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files]
                 edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic")
                 if st.button("Fine-Tune with Dataset 🔄"):
                     images = [Image.open(row["image"]) for _, row in edited_data.iterrows()]
@@ -664,12 +644,14 @@ 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']['sft'] = f"Fine-tuned Diffusion: {new_model_name}"
+                    entry = f"Fine-tuned Diffusion: {new_model_name}"
+                    if entry not in st.session_state['history']:
+                        st.session_state['history'].append(entry)
                     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:
                     writer = csv.writer(f)
-                    writer.writerow(["image", "text"])
+                    writer.writerow(["image", "text()])
                     for _, row in edited_data.iterrows():
                         writer.writerow([row["image"], row["text"]])
                 st.markdown(get_download_link(csv_path, "text/csv", "Download SFT Dataset CSV"), unsafe_allow_html=True)
@@ -696,7 +678,9 @@ 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']['test'] = f"Causal LM Test: {test_prompt} -> {result}"
+                entry = f"Causal LM Test: {test_prompt} -> {result}"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 st.write(f"**Generated Response**: {result}")
                 status_container.empty()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
@@ -705,8 +689,9 @@ with tab5:
                 image = st.session_state['builder'].generate(test_prompt)
                 output_file = generate_filename("diffusion_test", "png")
                 image.save(output_file)
-                st.session_state['captured_files']['diffusion_test'] = output_file
-                st.session_state['history']['test'] = f"Diffusion Test: {test_prompt} -> {output_file}"
+                entry = f"Diffusion Test: {test_prompt} -> {output_file}"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 st.image(image, caption="Generated Image")
                 update_gallery()
 
@@ -720,28 +705,31 @@ 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']['rag'] = f"NLP RAG Demo: Planned party at Wayne Manor"
+                entry = f"NLP RAG Demo: Planned party at Wayne Manor"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 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']['rag'] = f"CV RAG Demo: Generated party images"
+                entry = f"CV RAG Demo: Generated party images"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 st.dataframe(plan_df)
                 for _, row in plan_df.iterrows():
                     image = agent.generate(row["Image Idea"])
                     output_file = generate_filename(f"cv_rag_{row['Theme'].lower()}", "png")
                     image.save(output_file)
-                    st.session_state['captured_files'][f"cv_rag_{row['Theme'].lower()}"] = output_file
                     st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
                 update_gallery()
 
 with tab7:
     st.header("Test OCR 🔍")
-    captured_files = list(st.session_state['captured_files'].values())
+    captured_files = get_gallery_files(["png"])
     if captured_files:
-        selected_file = st.selectbox("Select Image", [f for f in captured_files if f and f.endswith(".png")], key="ocr_select")
+        selected_file = st.selectbox("Select Image", captured_files, key="ocr_select")
         if selected_file:
             image = Image.open(selected_file)
             st.image(image, caption="Input Image", use_container_width=True)
@@ -749,7 +737,9 @@ 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']['ocr'] = f"OCR Test: {selected_file} -> {output_file}"
+                entry = f"OCR Test: {selected_file} -> {output_file}"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 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
@@ -758,9 +748,9 @@ with tab7:
 
 with tab8:
     st.header("Test Image Gen 🎨")
-    captured_files = list(st.session_state['captured_files'].values())
+    captured_files = get_gallery_files(["png"])
     if captured_files:
-        selected_file = st.selectbox("Select Image", [f for f in captured_files if f and f.endswith(".png")], key="gen_select")
+        selected_file = st.selectbox("Select Image", captured_files, key="gen_select")
         if selected_file:
             image = Image.open(selected_file)
             st.image(image, caption="Reference Image", use_container_width=True)
@@ -769,7 +759,9 @@ 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']['gen'] = f"Image Gen Test: {prompt} -> {output_file}"
+                entry = f"Image Gen Test: {prompt} -> {output_file}"
+                if entry not in st.session_state['history']:
+                    st.session_state['history'].append(entry)
                 st.image(result, caption="Generated Image", use_container_width=True)
                 st.success(f"Image saved to {output_file}")
                 st.session_state['processing']['gen'] = False
@@ -779,10 +771,10 @@ with tab8:
 with tab9:
     st.header("Custom Diffusion 🎨🤓")
     st.write("Unleash your inner artist with our tiny diffusion models!")
-    captured_files = list(st.session_state['captured_files'].values())
+    captured_files = get_gallery_files(["png"])
     if captured_files:
         st.subheader("Select Images to Train")
-        selected_files = st.multiselect("Pick Images", [f for f in captured_files if f and f.endswith(".png")], key="diffusion_select")
+        selected_files = st.multiselect("Pick Images", captured_files, key="diffusion_select")
         images = [Image.open(file) for file in selected_files]
         
         model_options = [
@@ -803,8 +795,9 @@ with tab9:
                 builder.load_model(model_name)
                 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']['diffusion'] = f"Custom Diffusion: {model_choice} -> {output_file}"
+            entry = f"Custom Diffusion: {model_choice} -> {output_file}"
+            if entry not in st.session_state['history']:
+                st.session_state['history'].append(entry)
             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