Update app.py

app.py

@@ -49,6 +49,8 @@ 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
+if 'active_tab' not in st.session_state:
+    st.session_state['active_tab'] = "Build Titan 🌱"
 
 # Model Configuration Classes
 @dataclass
@@ -191,18 +193,19 @@ class DiffusionBuilder:
         dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
         optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
         self.pipeline.unet.train()
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         for epoch in range(epochs):
             with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"):
                 total_loss = 0
                 for batch in dataloader:
                     optimizer.zero_grad()
-                    image = batch["image"][0].to(self.pipeline.device)
+                    image = batch["image"][0].to(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(latant)
+                    latents = self.pipeline.vae.encode(torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float().to(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)
-                    text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
+                    text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(device))[0]
                     pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
                     loss = torch.nn.functional.mse_loss(pred_noise, noise)
                     loss.backward()
@@ -225,7 +228,7 @@ def generate_filename(sequence, ext="png"):
     import pytz
     central = pytz.timezone('US/Central')
     dt = datetime.now(central)
-    return f"{dt.strftime('%m-%d-%Y-%I-%M-%p')}.{ext}"
+    return f"{dt.strftime('%m-%d-%Y-%I-%M-%S-%p')}.{ext}"
 
 def get_download_link(file_path, mime_type="text/plain", label="Download"):
     with open(file_path, 'rb') as f:
@@ -244,8 +247,7 @@ def get_model_files(model_type="causal_lm"):
     return [d for d in glob.glob(path) if os.path.isdir(d)]
 
 def get_gallery_files(file_types):
-    files = sorted(list(set(f for ext in file_types for f in glob.glob(f"*.{ext}"))))  # Remove duplicates and sort
-    return files
+    return sorted(list(set(f for ext in file_types for f in glob.glob(f"*.{ext}"))))  # Remove duplicates and sort
 
 def update_gallery():
     media_files = get_gallery_files(["png"])
@@ -337,11 +339,19 @@ if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     st.rerun()
 
 # Tabs
-tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs([
+tabs = [
     "Build Titan 🌱", "Camera Snap 📷", 
     "Fine-Tune Titan (NLP) 🔧", "Test Titan (NLP) 🧪", "Agentic RAG Party (NLP) 🌐",
     "Fine-Tune Titan (CV) 🔧", "Test Titan (CV) 🧪", "Agentic RAG Party (CV) 🌐"
-])
+]
+tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs(tabs)
+
+# Log Tab Switches
+for i, tab in enumerate(tabs):
+    if st.session_state['active_tab'] != tab and st.session_state.get(f'tab{i}_active', False):
+        logger.info(f"Switched to tab: {tab}")
+        st.session_state['active_tab'] = tab
+    st.session_state[f'tab{i}_active'] = (st.session_state['active_tab'] == tab)
 
 with tab1:
     st.header("Build Titan 🌱")
@@ -350,9 +360,9 @@ with tab1:
         ["HuggingFaceTB/SmolLM-135M", "HuggingFaceTB/SmolLM-360M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else 
         ["stabilityai/stable-diffusion-2-base", "runwayml/stable-diffusion-v1-5"])
     model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
-    domain = st.text_input("Target Domain", "general", help="Where will your Titan flex its muscles? 💪")
+    domain = st.text_input("Target Domain", "general", help="Where will your Titan flex its muscles? 💪") if model_type == "Causal LM" else None
     if st.button("Download Model ⬇️"):
-        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model, size="small", domain=domain if model_type == "Causal LM" else None)
+        config = ModelConfig(name=model_name, base_model=base_model, size="small", domain=domain) if model_type == "Causal LM" else DiffusionConfig(name=model_name, base_model=base_model, size="small")
         builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
         builder.load_model(base_model, config)
         builder.save_model(config.model_path)