Update app.py

app.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import os
-import shutil
 import glob
 import base64
 import streamlit as st
@@ -31,7 +30,7 @@ st.set_page_config(
     initial_sidebar_state="expanded",
     menu_items={
         'Get Help': 'https://huggingface.co/awacke1',
-        'Report a bug': 'https://huggingface.co/spaces/awacke1',
+        'Report a Bug': 'https://huggingface.co/spaces/awacke1',
         'About': "Tiny Titans: Small models, big dreams, and a sprinkle of chaos! 🌌"
     }
 )
@@ -177,9 +176,9 @@ class DiffusionBuilder:
                 total_loss = 0
                 for batch in dataloader:
                     optimizer.zero_grad()
-                    image = batch["image"].to(self.pipeline.device)
-                    text = batch["text"]
-                    latents = self.pipeline.vae.encode(image).latent_dist.sample()
+                    image = batch["image"][0].to(self.pipeline.device)
+                    text = batch["text"][0]
+                    latents = self.pipeline.vae.encode(torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device)).latent_dist.sample()
                     noise = torch.randn_like(latents)
                     timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device)
                     noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
@@ -220,9 +219,10 @@ 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}")])
 
+# Mock Search Tool for RAG
 def mock_search(query: str) -> str:
     if "superhero" in query.lower():
-        return "Latest trends for 2025: Gold-plated Batman statues, VR superhero battles."
+        return "Latest trends: Gold-plated Batman statues, VR superhero battles."
     return "No relevant results found."
 
 class PartyPlannerAgent:
@@ -291,6 +291,7 @@ if media_files:
     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)
 
 st.sidebar.subheader("Model Management 🗂️")
 model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"])
@@ -351,7 +352,7 @@ with tab2:
                         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)  # Adjust frame rate
+                    time.sleep(1.0 / slice_count)
             st.session_state['captured_images'].extend(st.session_state['cam0_frames'])
             update_gallery()
             for frame in st.session_state['cam0_frames']:
@@ -379,7 +380,7 @@ with tab2:
                         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)  # Adjust frame rate
+                    time.sleep(1.0 / slice_count)
             st.session_state['captured_images'].extend(st.session_state['cam1_frames'])
             update_gallery()
             for frame in st.session_state['cam1_frames']:
@@ -420,6 +421,13 @@ with tab3:
                     zip_path = f"{new_config.model_path}.zip"
                     zip_directory(new_config.model_path, zip_path)
                     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"])
+                    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)
 
 with tab4:
     st.header("Test Titan 🧪")
@@ -456,4 +464,11 @@ with tab5:
                 st.dataframe(plan_df)
                 for _, row in plan_df.iterrows():
                     image = agent.generate(row["Image Idea"])
-                    st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
+                    st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
+
+# Main App
+st.sidebar.subheader("Action Logs 📜")
+log_container = st.sidebar.empty()
+with log_container:
+    for record in logger.handlers[0].buffer:
+        st.write(f"{record.asctime} - {record.levelname} - {record.message}")