Update app.py

app.py

@@ -45,10 +45,14 @@ st.set_page_config(
 # Initialize st.session_state
 if 'captured_images' not in st.session_state:
     st.session_state['captured_images'] = []
-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 'nlp_builder' not in st.session_state:
+    st.session_state['nlp_builder'] = None
+if 'cv_builder' not in st.session_state:
+    st.session_state['cv_builder'] = None
+if 'nlp_loaded' not in st.session_state:
+    st.session_state['nlp_loaded'] = False
+if 'cv_loaded' not in st.session_state:
+    st.session_state['cv_loaded'] = False
 if 'active_tab' not in st.session_state:
     st.session_state['active_tab'] = "Build Titan 🌱"
 
@@ -113,50 +117,76 @@ class ModelBuilder:
         self.sft_data = None
         self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! 😂", "Training complete! Time for a binary coffee break. ☕"]
     def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
-        with st.spinner(f"Loading {model_path}... ⏳ (Patience, young padawan!)"):
-            self.model = AutoModelForCausalLM.from_pretrained(model_path)
-            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            if config:
-                self.config = config
-            self.model.to("cuda" if torch.cuda.is_available() else "cpu")
-        st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
+        try:
+            with st.spinner(f"Loading {model_path}... ⏳ (Patience, young padawan!)"):
+                self.model = AutoModelForCausalLM.from_pretrained(model_path)
+                self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+                if self.tokenizer.pad_token is None:
+                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                if config:
+                    self.config = config
+                self.model.to("cuda" if torch.cuda.is_available() else "cpu")
+            st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
+            logger.info(f"Successfully loaded Causal LM model: {model_path}")
+        except torch.cuda.OutOfMemoryError as e:
+            st.error(f"GPU memory error loading {model_path}: {str(e)} 💥 (Out of GPU juice!)")
+            logger.error(f"GPU memory error loading {model_path}: {str(e)}")
+            raise
+        except MemoryError as e:
+            st.error(f"CPU memory error loading {model_path}: {str(e)} 💥 (RAM ran away!)")
+            logger.error(f"CPU memory error loading {model_path}: {str(e)}")
+            raise
+        except Exception as e:
+            st.error(f"Failed to load {model_path}: {str(e)} 💥 (Something broke—check the logs!)")
+            logger.error(f"Failed to load {model_path}: {str(e)}")
+            raise
         return self
     def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
-        self.sft_data = []
-        with open(csv_path, "r") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
-        dataset = SFTDataset(self.sft_data, self.tokenizer)
-        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
-        optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
-        self.model.train()
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(device)
-        for epoch in range(epochs):
-            with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️ (The AI is lifting weights!)"):
-                total_loss = 0
-                for batch in dataloader:
-                    optimizer.zero_grad()
-                    input_ids = batch["input_ids"].to(device)
-                    attention_mask = batch["attention_mask"].to(device)
-                    labels = batch["labels"].to(device)
-                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
-                    loss = outputs.loss
-                    loss.backward()
-                    optimizer.step()
-                    total_loss += loss.item()
-                st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
-        st.success(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}")
+        try:
+            self.sft_data = []
+            with open(csv_path, "r") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
+            dataset = SFTDataset(self.sft_data, self.tokenizer)
+            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+            optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
+            self.model.train()
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model.to(device)
+            for epoch in range(epochs):
+                with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️ (The AI is lifting weights!)"):
+                    total_loss = 0
+                    for batch in dataloader:
+                        optimizer.zero_grad()
+                        input_ids = batch["input_ids"].to(device)
+                        attention_mask = batch["attention_mask"].to(device)
+                        labels = batch["labels"].to(device)
+                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+                        loss = outputs.loss
+                        loss.backward()
+                        optimizer.step()
+                        total_loss += loss.item()
+                    st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
+            st.success(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}")
+            logger.info(f"Successfully fine-tuned Causal LM model: {self.config.name}")
+        except Exception as e:
+            st.error(f"Fine-tuning failed: {str(e)} 💥 (Training hit a snag!)")
+            logger.error(f"Fine-tuning failed: {str(e)}")
+            raise
         return self
     def save_model(self, path: str):
-        with st.spinner("Saving model... 💾 (Packing the AI’s suitcase!)"):
-            os.makedirs(os.path.dirname(path), exist_ok=True)
-            self.model.save_pretrained(path)
-            self.tokenizer.save_pretrained(path)
-        st.success(f"Model saved at {path}! ✅ May the force be with it.")
+        try:
+            with st.spinner("Saving model... 💾 (Packing the AI’s suitcase!)"):
+                os.makedirs(os.path.dirname(path), exist_ok=True)
+                self.model.save_pretrained(path)
+                self.tokenizer.save_pretrained(path)
+            st.success(f"Model saved at {path}! ✅ May the force be with it.")
+            logger.info(f"Model saved at {path}")
+        except Exception as e:
+            st.error(f"Failed to save model: {str(e)} 💥 (Save operation crashed!)")
+            logger.error(f"Failed to save model: {str(e)}")
+            raise
     def evaluate(self, prompt: str, status_container=None):
         self.model.eval()
         if status_container:
@@ -181,46 +211,77 @@ class DiffusionBuilder:
         self.pipeline = None
     def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
         from diffusers import StableDiffusionPipeline
-        with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
-            self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
-            self.pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
-            if config:
-                self.config = config
-        st.success(f"Diffusion model loaded! 🎨")
+        try:
+            with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
+                self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
+                self.pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+                if config:
+                    self.config = config
+            st.success(f"Diffusion model loaded! 🎨")
+            logger.info(f"Successfully loaded Diffusion model: {model_path}")
+        except torch.cuda.OutOfMemoryError as e:
+            st.error(f"GPU memory error loading {model_path}: {str(e)} 💥 (Out of GPU juice!)")
+            logger.error(f"GPU memory error loading {model_path}: {str(e)}")
+            raise
+        except MemoryError as e:
+            st.error(f"CPU memory error loading {model_path}: {str(e)} 💥 (RAM ran away!)")
+            logger.error(f"CPU memory error loading {model_path}: {str(e)}")
+            raise
+        except Exception as e:
+            st.error(f"Failed to load {model_path}: {str(e)} 💥 (Something broke—check the logs!)")
+            logger.error(f"Failed to load {model_path}: {str(e)}")
+            raise
         return self
     def fine_tune_sft(self, images, texts, epochs=3):
-        dataset = DiffusionDataset(images, texts)
-        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(device)
-                    text = batch["text"][0]
-                    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(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()
-                    optimizer.step()
-                    total_loss += loss.item()
-                st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
-        st.success("Diffusion SFT Fine-tuning completed! 🎨")
+        try:
+            dataset = DiffusionDataset(images, texts)
+            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(device)
+                        text = batch["text"][0]
+                        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(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()
+                        optimizer.step()
+                        total_loss += loss.item()
+                    st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
+            st.success("Diffusion SFT Fine-tuning completed! 🎨")
+            logger.info(f"Successfully fine-tuned Diffusion model: {self.config.name}")
+        except Exception as e:
+            st.error(f"Fine-tuning failed: {str(e)} 💥 (Training hit a snag!)")
+            logger.error(f"Fine-tuning failed: {str(e)}")
+            raise
         return self
     def save_model(self, path: str):
-        with st.spinner("Saving diffusion model... 💾"):
-            os.makedirs(os.path.dirname(path), exist_ok=True)
-            self.pipeline.save_pretrained(path)
-        st.success(f"Diffusion model saved at {path}! ✅")
+        try:
+            with st.spinner("Saving diffusion model... 💾"):
+                os.makedirs(os.path.dirname(path), exist_ok=True)
+                self.pipeline.save_pretrained(path)
+            st.success(f"Diffusion model saved at {path}! ✅")
+            logger.info(f"Diffusion model saved at {path}")
+        except Exception as e:
+            st.error(f"Failed to save model: {str(e)} 💥 (Save operation crashed!)")
+            logger.error(f"Failed to save model: {str(e)}")
+            raise
     def generate(self, prompt: str):
-        return self.pipeline(prompt, num_inference_steps=50).images[0]
+        try:
+            return self.pipeline(prompt, num_inference_steps=50).images[0]
+        except Exception as e:
+            st.error(f"Image generation failed: {str(e)} 💥 (Pixel party pooper!)")
+            logger.error(f"Image generation failed: {str(e)}")
+            raise
 
 # Utility Functions
 def generate_filename(sequence, ext="png"):
@@ -231,16 +292,24 @@ def generate_filename(sequence, ext="png"):
     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:
-        data = f.read()
-    b64 = base64.b64encode(data).decode()
-    return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
+    try:
+        with open(file_path, 'rb') as f:
+            data = f.read()
+        b64 = base64.b64encode(data).decode()
+        return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
+    except Exception as e:
+        logger.error(f"Failed to generate download link for {file_path}: {str(e)}")
+        return f"Error: Could not generate link for {file_path}"
 
 def zip_directory(directory_path, zip_path):
-    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(directory_path):
-            for file in files:
-                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))
+    try:
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, _, files in os.walk(directory_path):
+                for file in files:
+                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))
+    except Exception as e:
+        logger.error(f"Failed to zip directory {directory_path}: {str(e)}")
+        raise
 
 def get_model_files(model_type="causal_lm"):
     path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
@@ -321,7 +390,7 @@ def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_
 # Main App
 st.title("SFT Tiny Titans 🚀 (Small but Mighty!)")
 
-# Sidebar Galleries
+# Sidebar Galleries and Status
 st.sidebar.header("Media Gallery 🎨")
 gallery_size = st.sidebar.slider("Gallery Size 📸", 1, 10, 4, help="Adjust how many epic captures you see! 🌟")
 update_gallery()
@@ -333,10 +402,21 @@ selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dir
 if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
     config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small")
-    builder.load_model(selected_model, config)
-    st.session_state['builder'] = builder
-    st.session_state['model_loaded'] = True
-    st.rerun()
+    try:
+        builder.load_model(selected_model, config)
+        if model_type == "Causal LM":
+            st.session_state['nlp_builder'] = builder
+            st.session_state['nlp_loaded'] = True
+        else:
+            st.session_state['cv_builder'] = builder
+            st.session_state['cv_loaded'] = True
+        st.rerun()
+    except Exception as e:
+        st.error(f"Model load failed: {str(e)} 💥 (Check logs for details!)")
+
+st.sidebar.subheader("Model Status 🚦")
+st.sidebar.write(f"**NLP Model**: {'Loaded' if st.session_state['nlp_loaded'] else 'Not Loaded'} {'(Active)' if st.session_state['nlp_loaded'] and isinstance(st.session_state.get('nlp_builder'), ModelBuilder) else ''}")
+st.sidebar.write(f"**CV Model**: {'Loaded' if st.session_state['cv_loaded'] else 'Not Loaded'} {'(Active)' if st.session_state['cv_loaded'] and isinstance(st.session_state.get('cv_builder'), DiffusionBuilder) else ''}")
 
 # Tabs
 tabs = [
@@ -364,11 +444,18 @@ with tab1:
     if st.button("Download Model ⬇️"):
         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)
-        st.session_state['builder'] = builder
-        st.session_state['model_loaded'] = True
-        st.rerun()
+        try:
+            builder.load_model(base_model, config)
+            builder.save_model(config.model_path)
+            if model_type == "Causal LM":
+                st.session_state['nlp_builder'] = builder
+                st.session_state['nlp_loaded'] = True
+            else:
+                st.session_state['cv_builder'] = builder
+                st.session_state['cv_loaded'] = True
+            st.rerun()
+        except Exception as e:
+            st.error(f"Model build failed: {str(e)} 💥 (Check logs for details!)")
 
 with tab2:
     st.header("Camera Snap 📷 (Dual Capture!)")
@@ -402,7 +489,7 @@ with tab2:
 
 with tab3:  # Fine-Tune Titan (NLP)
     st.header("Fine-Tune Titan (NLP) 🔧 (Teach Your Word Wizard Some Tricks!)")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], ModelBuilder):
+    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
         st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no magic!)")
     else:
         if st.button("Generate Sample CSV 📝"):
@@ -423,12 +510,12 @@ with tab3:  # Fine-Tune Titan (NLP)
             csv_path = f"uploaded_sft_data_{int(time.time())}.csv"
             with open(csv_path, "wb") as f:
                 f.write(uploaded_csv.read())
-            new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
-            new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small", domain=st.session_state['builder'].config.domain)
-            st.session_state['builder'].config = new_config
+            new_model_name = f"{st.session_state['nlp_builder'].config.name}-sft-{int(time.time())}"
+            new_config = ModelConfig(name=new_model_name, base_model=st.session_state['nlp_builder'].config.base_model, size="small", domain=st.session_state['nlp_builder'].config.domain)
+            st.session_state['nlp_builder'].config = new_config
             with st.status("Fine-tuning NLP Titan... ⏳ (Whipping words into shape!)", expanded=True) as status:
-                st.session_state['builder'].fine_tune_sft(csv_path)
-                st.session_state['builder'].save_model(new_config.model_path)
+                st.session_state['nlp_builder'].fine_tune_sft(csv_path)
+                st.session_state['nlp_builder'].save_model(new_config.model_path)
                 status.update(label="Fine-tuning completed! 🎉 (Wordsmith Titan unleashed!)", state="complete")
             zip_path = f"{new_config.model_path}.zip"
             zip_directory(new_config.model_path, zip_path)
@@ -436,17 +523,17 @@ with tab3:  # Fine-Tune Titan (NLP)
 
 with tab4:  # Test Titan (NLP)
     st.header("Test Titan (NLP) 🧪 (Put Your Word Wizard to the Test!)")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], ModelBuilder):
+    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
         st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no test drive!)")
     else:
-        if st.session_state['builder'].sft_data:
+        if st.session_state['nlp_builder'].sft_data:
             st.write("Testing with SFT Data:")
             with st.spinner("Running SFT data tests... ⏳ (Titan’s flexing its word muscles!)"):
-                for item in st.session_state['builder'].sft_data[:3]:
+                for item in st.session_state['nlp_builder'].sft_data[:3]:
                     prompt = item["prompt"]
                     expected = item["response"]
                     status_container = st.empty()
-                    generated = st.session_state['builder'].evaluate(prompt, status_container)
+                    generated = st.session_state['nlp_builder'].evaluate(prompt, status_container)
                     st.write(f"**Prompt**: {prompt}")
                     st.write(f"**Expected**: {expected}")
                     st.write(f"**Generated**: {generated} (Titan says: '{random.choice(['Bleep bloop!', 'I am groot!', '42!'])}')")
@@ -456,19 +543,19 @@ with tab4:  # Test Titan (NLP)
         if st.button("Run Test ▶️"):
             with st.spinner("Testing your prompt... ⏳ (Titan’s pondering deeply!)"):
                 status_container = st.empty()
-                result = st.session_state['builder'].evaluate(test_prompt, status_container)
+                result = st.session_state['nlp_builder'].evaluate(test_prompt, status_container)
                 st.write(f"**Generated Response**: {result} (Titan’s wisdom unleashed!)")
                 status_container.empty()
 
 with tab5:  # Agentic RAG Party (NLP)
     st.header("Agentic RAG Party (NLP) 🌐 (Party Like It’s 2099!)")
     st.write("This demo uses your SFT-tuned NLP Titan to plan a superhero party with mock retrieval!")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], ModelBuilder):
+    if not st.session_state['nlp_loaded'] or not isinstance(st.session_state['nlp_builder'], ModelBuilder):
         st.warning("Please build or load an NLP Titan first! ⚠️ (No word wizard, no party!)")
     else:
         if st.button("Run NLP RAG Demo 🎉"):
             with st.spinner("Loading your SFT-tuned NLP Titan... ⏳ (Titan’s suiting up!)"):
-                agent = PartyPlannerAgent(st.session_state['builder'].model, st.session_state['builder'].tokenizer)
+                agent = PartyPlannerAgent(st.session_state['nlp_builder'].model, st.session_state['nlp_builder'].tokenizer)
                 st.write("Agent ready! 🦸‍♂️ (Time to plan an epic bash!)")
             task = """
             Plan a luxury superhero-themed party at Wayne Manor (42.3601° N, 71.0589° W). 
@@ -508,10 +595,11 @@ with tab5:  # Agentic RAG Party (NLP)
                     st.write("Party on, Wayne! 🦸‍♂️🎉")
                 except Exception as e:
                     st.error(f"Error planning party: {str(e)} (Even Superman has kryptonite days!)")
+                    logger.error(f"Error in NLP RAG demo: {str(e)}")
 
 with tab6:  # Fine-Tune Titan (CV)
     st.header("Fine-Tune Titan (CV) 🔧 (Paint Your Titan’s Masterpiece!)")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], DiffusionBuilder):
+    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
         st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no canvas!)")
     else:
         captured_images = get_gallery_files(["png"])
@@ -521,12 +609,12 @@ with tab6:  # Fine-Tune Titan (CV)
             if st.button("Fine-Tune with Dataset 🔄"):
                 images = [Image.open(row["image"]) for _, row in edited_data.iterrows()]
                 texts = [row["text"] for _, row in edited_data.iterrows()]
-                new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
-                new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
-                st.session_state['builder'].config = new_config
+                new_model_name = f"{st.session_state['cv_builder'].config.name}-sft-{int(time.time())}"
+                new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['cv_builder'].config.base_model, size="small")
+                st.session_state['cv_builder'].config = new_config
                 with st.status("Fine-tuning CV Titan... ⏳ (Brushing up those pixels!)", expanded=True) as status:
-                    st.session_state['builder'].fine_tune_sft(images, texts)
-                    st.session_state['builder'].save_model(new_config.model_path)
+                    st.session_state['cv_builder'].fine_tune_sft(images, texts)
+                    st.session_state['cv_builder'].save_model(new_config.model_path)
                     status.update(label="Fine-tuning completed! 🎉 (Pixel Titan unleashed!)", state="complete")
                 zip_path = f"{new_config.model_path}.zip"
                 zip_directory(new_config.model_path, zip_path)
@@ -541,32 +629,36 @@ with tab6:  # Fine-Tune Titan (CV)
 
 with tab7:  # Test Titan (CV)
     st.header("Test Titan (CV) 🧪 (Unleash Your Pixel Power!)")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], DiffusionBuilder):
+    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
         st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no masterpiece!)")
     else:
         test_prompt = st.text_area("Enter Test Prompt 🎨", "Neon Batman", help="Dream up a wild image—your Titan’s got the brush! 🖌️")
         if st.button("Run Test ▶️"):
             with st.spinner("Painting your masterpiece... ⏳ (Titan’s mixing colors!)"):
-                image = st.session_state['builder'].generate(test_prompt)
+                image = st.session_state['cv_builder'].generate(test_prompt)
                 st.image(image, caption="Generated Image", use_container_width=True)
 
 with tab8:  # Agentic RAG Party (CV)
     st.header("Agentic RAG Party (CV) 🌐 (Party with Pixels!)")
     st.write("This demo uses your SFT-tuned CV Titan to generate superhero party images with mock retrieval!")
-    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], DiffusionBuilder):
+    if not st.session_state['cv_loaded'] or not isinstance(st.session_state['cv_builder'], DiffusionBuilder):
         st.warning("Please build or load a CV Titan first! ⚠️ (No artist, no party!)")
     else:
         if st.button("Run CV RAG Demo 🎉"):
             with st.spinner("Loading your SFT-tuned CV Titan... ⏳ (Titan’s grabbing its paintbrush!)"):
-                agent = CVPartyPlannerAgent(st.session_state['builder'].pipeline)
+                agent = CVPartyPlannerAgent(st.session_state['cv_builder'].pipeline)
                 st.write("Agent ready! 🎨 (Time to paint an epic bash!)")
             task = "Generate images for a luxury superhero-themed party."
             with st.spinner("Crafting superhero party visuals... ⏳ (Pixels assemble!)"):
-                plan_df = agent.plan_party(task)
-                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']}", use_container_width=True)
+                try:
+                    plan_df = agent.plan_party(task)
+                    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']}", use_container_width=True)
+                except Exception as e:
+                    st.error(f"Error in CV RAG demo: {str(e)} 💥 (Pixel party crashed!)")
+                    logger.error(f"Error in CV RAG demo: {str(e)}")
 
 # Display Logs
 st.sidebar.subheader("Action Logs 📜")