Update app.py

app.py

@@ -8,9 +8,16 @@ from dataclasses import dataclass
 import zipfile
 import logging
 
-# Logging setup
+# Logging setup with custom log storage
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
+log_records = []  # Custom list to store logs
+
+class LogCaptureHandler(logging.Handler):
+    def emit(self, record):
+        log_records.append(record)
+
+logger.addHandler(LogCaptureHandler())
 
 st.set_page_config(page_title="SFT Tiny Titans 🚀", page_icon="🤖", layout="wide", initial_sidebar_state="expanded")
 
@@ -132,7 +139,7 @@ class DiffusionBuilder:
             for _ in range(1):
                 for img, text in zip(images, texts):
                     optimizer.zero_grad()
-                    img_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device) / 255.0  # Normalize
+                    img_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device) / 255.0
                     latents = self.pipeline.vae.encode(img_tensor).latent_dist.sample()
                     noise = torch.randn_like(latents)
                     timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (1,), device=latents.device)
@@ -368,5 +375,5 @@ with tab4:
 st.sidebar.subheader("Action Logs 📜")
 log_container = st.sidebar.empty()
 with log_container:
-    for record in logger.handlers[0].buffer:
+    for record in log_records:
         st.write(f"{record.asctime} - {record.levelname} - {record.message}")