app.py

#!/usr/bin/env python3
import os
import glob
import base64
import time
import shutil
import streamlit as st
import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
from diffusers import StableDiffusionPipeline
from torch.utils.data import Dataset, DataLoader
import csv
import fitz  # PyMuPDF
import requests
from PIL import Image
import cv2
import numpy as np
import logging
import asyncio
import aiofiles
from io import BytesIO
from dataclasses import dataclass
from typing import Optional, Tuple
import zipfile
import math
import random

# Logging setup with custom buffer
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
log_records = []

class LogCaptureHandler(logging.Handler):
    def emit(self, record):
        log_records.append(record)

logger.addHandler(LogCaptureHandler())

# Page Configuration
st.set_page_config(
    page_title="AI Vision & SFT Titans 🚀",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://huggingface.co/awacke1',
        'Report a Bug': 'https://huggingface.co/spaces/awacke1',
        'About': "AI Vision & SFT Titans: PDFs, OCR, Image Gen, Line Drawings, and SFT on CPU! 🌌"
    }
)

# Initialize st.session_state
if 'captured_files' not in st.session_state:
    st.session_state['captured_files'] = []
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 'processing' not in st.session_state:
    st.session_state['processing'] = {}
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Model Configuration Classes
@dataclass
class ModelConfig:
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    model_type: str = "causal_lm"
    @property
    def model_path(self):
        return f"models/{self.name}"

@dataclass
class DiffusionConfig:
    name: str
    base_model: str
    size: str
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

# Datasets
class SFTDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        prompt = self.data[idx]["prompt"]
        response = self.data[idx]["response"]
        full_text = f"{prompt} {response}"
        full_encoding = self.tokenizer(full_text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt")
        prompt_encoding = self.tokenizer(prompt, max_length=self.max_length, padding=False, truncation=True, return_tensors="pt")
        input_ids = full_encoding["input_ids"].squeeze()
        attention_mask = full_encoding["attention_mask"].squeeze()
        labels = input_ids.clone()
        prompt_len = prompt_encoding["input_ids"].shape[1]
        if prompt_len < self.max_length:
            labels[:prompt_len] = -100
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

class DiffusionDataset(Dataset):
    def __init__(self, images, texts):
        self.images = images
        self.texts = texts
    def __len__(self):
        return len(self.images)
    def __getitem__(self, idx):
        return {"image": self.images[idx], "text": self.texts[idx]}

# Model Builders
class ModelBuilder:
    def __init__(self):
        self.config = None
        self.model = None
        self.tokenizer = None
        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}... ⏳"):
            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)}")
        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}... ⚙️"):
                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)}")
        return self
    def save_model(self, path: str):
        with st.spinner("Saving model... 💾"):
            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}! ✅")
    def evaluate(self, prompt: str, status_container=None):
        self.model.eval()
        if status_container:
            status_container.write("Preparing to evaluate... 🧠")
        try:
            with torch.no_grad():
                inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
                outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
                return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            if status_container:
                status_container.error(f"Oops! Something broke: {str(e)} 💥")
            return f"Error: {str(e)}"

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
            self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
            if config:
                self.config = config
        st.success(f"Diffusion model loaded! 🎨")
        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()
        for epoch in range(epochs):
            with—for 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)
                    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)
                    text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.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! 🎨")
        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}! ✅")
    def generate(self, prompt: str):
        return self.pipeline(prompt, num_inference_steps=20).images[0]

# Utility Functions
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>'

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)))

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):
    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():
        return "Latest trends: Gold-plated Batman statues, VR superhero battles."
    return "No relevant results found."

def mock_duckduckgo_search(query: str) -> str:
    if "superhero party trends" in query.lower():
        return """
        Latest trends for 2025: 
        - Luxury decorations: Gold-plated Batman statues, holographic Avengers displays.
        - Entertainment: Live stunt shows with Iron Man suits, VR superhero battles.
        - Catering: Gourmet kryptonite-green cocktails, Thor’s hammer-shaped appetizers.
        """
    return "No relevant results found."

# Agent Classes
class PartyPlannerAgent:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
    def generate(self, prompt: str) -> str:
        self.model.eval()
        with torch.no_grad():
            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device)
            outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7)
            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_duckduckgo_search("latest superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        plan_text = self.generate(prompt)
        locations = {
            "Wayne Manor": (42.3601, -71.0589),
            "New York": (40.7128, -74.0060),
            "Los Angeles": (34.0522, -118.2437),
            "London": (51.5074, -0.1278)
        }
        wayne_coords = locations["Wayne Manor"]
        travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"}
        catchphrases = ["To the Batmobile!", "Avengers, assemble!", "I am Iron Man!", "By the power of Grayskull!"]
        data = [
            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Holographic Avengers displays", "Catchphrase": random.choice(catchphrases)},
            {"Location": "London", "Travel Time (hrs)": travel_times["London"], "Luxury Idea": "Live stunt shows with Iron Man suits", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles", "Catchphrase": random.choice(catchphrases)},
            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gourmet kryptonite-green cocktails", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Thor’s hammer-shaped appetizers", "Catchphrase": random.choice(catchphrases)},
        ]
        return pd.DataFrame(data)

class CVPartyPlannerAgent:
    def __init__(self, pipeline):
        self.pipeline = pipeline
    def generate(self, prompt: str) -> Image.Image:
        return self.pipeline(prompt, num_inference_steps=20).images[0]
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_search("superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        data = [
            {"Theme": "Batman", "Image Idea": "Gold-plated Batman statue"},
            {"Theme": "Avengers", "Image Idea": "VR superhero battle scene"}
        ]
        return pd.DataFrame(data)

def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_coords: Tuple[float, float], cruising_speed_kmh: float = 750.0) -> float:
    def to_radians(degrees: float) -> float:
        return degrees * (math.pi / 180)
    lat1, lon1 = map(to_radians, origin_coords)
    lat2, lon2 = map(to_radians, destination_coords)
    EARTH_RADIUS_KM = 6371.0
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = (math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2)
    c = 2 * math.asin(math.sqrt(a))
    distance = EARTH_RADIUS_KM * c
    actual_distance = distance * 1.1
    flight_time = (actual_distance / cruising_speed_kmh) + 1.0
    return round(flight_time, 2)

# Async Processing Functions
async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
    start_time = time.time()
    status = st.empty()
    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
    doc = fitz.open(pdf_path)
    output_files = []
    if mode == "thumbnail":
        page = doc[0]
        pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
        output_file = generate_filename("thumbnail", "png")
        pix.save(output_file)
        output_files.append(output_file)
    elif mode == "twopage":
        for i in range(min(2, len(doc))):
            page = doc[i]
            pix = page.get_pixmap(matrix=fitz.Matrix(1.0, 1.0))
            output_file = generate_filename(f"twopage_{i}", "png")
            pix.save(output_file)
            output_files.append(output_file)
    doc.close()
    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']:
            st.session_state['captured_files'].append(file)
    update_gallery()
    return output_files

async def process_ocr(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing GOT-OCR2_0... (0s)")
    tokenizer, model = load_ocr_got()
    result = model.chat(tokenizer, image, ocr_type='ocr')
    elapsed = int(time.time() - start_time)
    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']:
        st.session_state['captured_files'].append(output_file)
    update_gallery()
    return result

# Main App
st.title("AI Vision & SFT Titans 🚀")

# Sidebar
st.sidebar.header("Captured Files 📜")
gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
update_gallery()

st.sidebar.subheader("Model Management 🗂️")
model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"], key="sidebar_model_type")
model_dirs = get_model_files(model_type)
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs, key="sidebar_model_select")
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()

st.sidebar.subheader("Action Logs 📜")
log_container = st.sidebar.empty()
with log_container:
    for record in log_records:
        st.write(f"{record.asctime} - {record.levelname} - {record.message}")

st.sidebar.subheader("History 📜")
history_container = st.sidebar.empty()
with history_container:
    for entry in st.session_state['history'][-5:]:
        st.write(entry)

# Tabs
tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs([
    "Camera Snap 📷", "Download PDFs 📥", "Build Titan 🌱", "Fine-Tune Titan 🔧", 
    "Test Titan 🧪", "Agentic RAG Party 🌐", "Test OCR 🔍", "Test Image Gen 🎨"
])

with tab1:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
        if cam0_img:
            filename = generate_filename(0)
            if filename not in st.session_state['captured_files']:
                with open(filename, "wb") as f:
                    f.write(cam0_img.getvalue())
                st.image(Image.open(filename), caption=filename, use_container_width=True)
                logger.info(f"Saved snapshot from Camera 0: {filename}")
                st.session_state['captured_files'].append(filename)
                st.session_state['history'].append(f"Snapshot from Cam 0: {filename}")
                update_gallery()
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
        if cam1_img:
            filename = generate_filename(1)
            if filename not in st.session_state['captured_files']:
                with open(filename, "wb") as f:
                    f.write(cam1_img.getvalue())
                st.image(Image.open(filename), caption=filename, use_container_width=True)
                logger.info(f"Saved snapshot from Camera 1: {filename}")
                st.session_state['captured_files'].append(filename)
                st.session_state['history'].append(f"Snapshot from Cam 1: {filename}")
                update_gallery()

    st.subheader("Burst Capture")
    slice_count = st.number_input("Number of Frames", min_value=1, max_value=20, value=10, key="burst_count")
    if st.button("Start Burst Capture 📸"):
        st.session_state['burst_frames'] = []
        placeholder = st.empty()
        for i in range(slice_count):
            with placeholder.container():
                st.write(f"Capturing frame {i+1}/{slice_count}...")
                img = st.camera_input(f"Frame {i}", key=f"burst_{i}_{time.time()}")
                if img:
                    filename = generate_filename(f"burst_{i}")
                    if filename not in st.session_state['captured_files']:
                        with open(filename, "wb") as f:
                            f.write(img.getvalue())
                        st.session_state['burst_frames'].append(filename)
                        logger.info(f"Saved burst frame {i}: {filename}")
                        st.session_state['history'].append(f"Burst frame {i}: {filename}")
                        st.image(Image.open(filename), caption=filename, use_container_width=True)
                    time.sleep(0.5)
        st.session_state['captured_files'].extend([f for f in st.session_state['burst_frames'] if f not in st.session_state['captured_files']])
        update_gallery()
        placeholder.success(f"Captured {len(st.session_state['burst_frames'])} frames!")

with tab2:
    st.header("Download PDFs 📥")
    url_input = st.text_area("Enter PDF URLs (one per line)", height=100)
    mode = st.selectbox("Snapshot Mode", ["Thumbnail", "Two-Page View"], key="download_mode")
    if st.button("Download & Snapshot 📸"):
        urls = url_input.strip().split("\n")
        for url in urls:
            if url:
                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'].append(f"Downloaded PDF: {pdf_path}")
                    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)
                else:
                    st.error(f"Failed to download {url}")

with tab3:
    st.header("Build Titan 🌱")
    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
    base_model = st.selectbox("Select Tiny Model", 
        ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else 
        ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"])
    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
    domain = st.text_input("Target Domain", "general")
    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)
        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.session_state['history'].append(f"Built {model_type} model: {model_name}")
        st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
        st.rerun()

with tab4:
    st.header("Fine-Tune Titan 🔧")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.button("Generate Sample CSV 📝"):
                sample_data = [
                    {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human smarts in machines."},
                    {"prompt": "Explain machine learning", "response": "Machine learning is AI’s gym where models bulk up on data."},
                ]
                csv_path = f"sft_data_{int(time.time())}.csv"
                with open(csv_path, "w", newline="") as f:
                    writer = csv.DictWriter(f, fieldnames=["prompt", "response"])
                    writer.writeheader()
                    writer.writerows(sample_data)
                st.markdown(get_download_link(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True)
                st.success(f"Sample CSV generated as {csv_path}! ✅")

            uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv")
            if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"):
                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
                st.session_state['builder'].fine_tune_sft(csv_path)
                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'].append(f"Fine-tuned Causal LM: {new_model_name}")
                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 = 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[:min(len(captured_files), slice_count)]]
                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()]
                    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
                    st.session_state['builder'].fine_tune_sft(images, texts)
                    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'].append(f"Fine-tuned Diffusion: {new_model_name}")
                    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 tab5:
    st.header("Test Titan 🧪")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.session_state['builder'].sft_data:
                st.write("Testing with SFT Data:")
                for item in st.session_state['builder'].sft_data[:3]:
                    prompt = item["prompt"]
                    expected = item["response"]
                    status_container = st.empty()
                    generated = st.session_state['builder'].evaluate(prompt, status_container)
                    st.write(f"**Prompt**: {prompt}")
                    st.write(f"**Expected**: {expected}")
                    st.write(f"**Generated**: {generated}")
                    st.write("---")
                    status_container.empty()
            test_prompt = st.text_area("Enter Test Prompt", "What is AI?")
            if st.button("Run Test ▶️"):
                status_container = st.empty()
                result = st.session_state['builder'].evaluate(test_prompt, status_container)
                st.session_state['history'].append(f"Causal LM Test: {test_prompt} -> {result}")
                st.write(f"**Generated Response**: {result}")
                status_container.empty()
        elif isinstance(st.session_state['builder'], DiffusionBuilder):
            test_prompt = st.text_area("Enter Test Prompt", "Neon Batman")
            if st.button("Run Test ▶️"):
                image = st.session_state['builder'].generate(test_prompt)
                output_file = generate_filename("diffusion_test", "png")
                image.save(output_file)
                st.session_state['captured_files'].append(output_file)
                st.session_state['history'].append(f"Diffusion Test: {test_prompt} -> {output_file}")
                st.image(image, caption="Generated Image")
                update_gallery()

with tab6:
    st.header("Agentic RAG Party 🌐")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.button("Run NLP RAG Demo 🎉"):
                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'].append(f"NLP RAG Demo: Planned party at Wayne Manor")
                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'].append(f"CV RAG Demo: Generated party images")
                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'].append(output_file)
                    st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
                update_gallery()

with tab7:
    st.header("Test OCR 🔍")
    captured_files = get_gallery_files(["png"])
    if captured_files:
        selected_file = st.selectbox("Select Image", captured_files, key="ocr_select")
        image = Image.open(selected_file)
        st.image(image, caption="Input Image", use_container_width=True)
        if st.button("Run OCR 🚀", key="ocr_run"):
            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'].append(f"OCR Test: {selected_file} -> {output_file}")
            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
    else:
        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")

with tab8:
    st.header("Test Image Gen 🎨")
    captured_files = get_gallery_files(["png"])
    if captured_files:
        selected_file = st.selectbox("Select Image", captured_files, key="gen_select")
        image = Image.open(selected_file)
        st.image(image, caption="Reference Image", use_container_width=True)
        prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
        if st.button("Run Image Gen 🚀", key="gen_run"):
            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'].append(f"Image Gen Test: {prompt} -> {output_file}")
            st.image(result, caption="Generated Image", use_container_width=True)
            st.success(f"Image saved to {output_file}")
            st.session_state['processing']['gen'] = False
    else:
        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")

# Initial Gallery Update
update_gallery()