Update app.py

app.py

@@ -1,25 +1,16 @@
 #!/usr/bin/env python3
 import os
-import glob
-import base64
 import streamlit as st
-import pandas as pd
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from torch.utils.data import Dataset, DataLoader
-import csv
-import time
-from dataclasses import dataclass
-from typing import Optional, Tuple
-import zipfile
-import math
 from PIL import Image
-import random
-import logging
+import torch
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, AutoTokenizer, AutoModel
+from diffusers import StableDiffusionPipeline
+import cv2
 import numpy as np
-from diffusers import StableDiffusionPipeline, DDPMPipeline, EulerAncestralDiscreteScheduler
+import logging
+from io import BytesIO
 
-# Logging setup with a custom buffer
+# Logging setup
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 log_records = []
@@ -32,196 +23,16 @@ logger.addHandler(LogCaptureHandler())
 
 # Page Configuration
 st.set_page_config(
-    page_title="SFT Tiny Titans 🚀",
+    page_title="AI Vision 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': "Tiny Titans: Small models, big dreams, and a sprinkle of chaos! 🌌"
-    }
+    menu_items={'About': "AI Vision Titans: OCR, Image Gen, Line Drawings on CPU! 🌌"}
 )
 
 # 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
-
-# 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
-        self.model_type = None
-    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None, model_type: str = "StableDiffusion", download: bool = True):
-        with st.spinner(f"{'Downloading' if download else 'Loading'} {model_path}... ⏳"):
-            if model_type == "StableDiffusion":
-                self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32, local_files_only=not download).to("cpu")
-            elif model_type == "DDPM":
-                self.pipeline = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch.float32, local_files_only=not download).to("cpu")
-                self.pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipeline.scheduler.config)
-            if config:
-                self.config = config
-            self.model_type = model_type
-        st.success(f"Diffusion model {'downloaded' if download else '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 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, image=None):
-        if self.model_type == "StableDiffusion":
-            return self.pipeline(prompt, num_inference_steps=50).images[0]
-        elif self.model_type == "DDPM":
-            return self.pipeline(num_inference_steps=50).images[0]
 
 # Utility Functions
 def generate_filename(sequence, ext="png"):
@@ -231,22 +42,6 @@ def generate_filename(sequence, ext="png"):
     timestamp = datetime.now(central).strftime("%d%m%Y%H%M%S%p")
     return f"{sequence}{timestamp}.{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>'
-
-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}")])
 
@@ -257,119 +52,55 @@ def update_gallery():
         for idx, file in enumerate(media_files[:gallery_size * 2]):
             with cols[idx % 2]:
                 st.image(Image.open(file), caption=file, use_container_width=True)
-                st.markdown(get_download_link(file, "image/png", "Download Image"), unsafe_allow_html=True)
-
-# 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."
 
-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_search("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)}
-        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"}
-        data = [
-            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues"},
-            {"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles"}
-        ]
-        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=50).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)
+# Model Loaders (Simplified, CPU-focused)
+def load_ocr_qwen2vl():
+    model_id = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
+    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    model = Qwen2VLForConditionalGeneration.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
+    return processor, model
+
+def load_ocr_got():
+    model_id = "ucaslcl/GOT-OCR2_0"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
+    return tokenizer, model
+
+def load_image_gen():
+    model_id = "OFA-Sys/small-stable-diffusion-v0"  # Small, CPU-friendly
+    pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32).to("cpu")
+    return pipeline
+
+def load_line_drawer():
+    # Simplified from your Torch Space (assuming a UNet-like model for edge detection)
+    # Placeholder: Using OpenCV edge detection as a minimal CPU example
+    def edge_detection(image):
+        img_np = np.array(image.convert("RGB"))
+        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
+        edges = cv2.Canny(gray, 100, 200)
+        return Image.fromarray(edges)
+    return edge_detection
 
 # Main App
-st.title("SFT Tiny Titans 🚀 (Small but Mighty!)")
+st.title("AI Vision Titans 🚀 (OCR, Gen, Drawings!)")
 
-# Sidebar Galleries
-st.sidebar.header("Media Gallery 🎨")
+# Sidebar Gallery
+st.sidebar.header("Captured Images 🎨")
 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"])
-model_dirs = get_model_files("causal_lm" if model_type == "Causal LM" else "diffusion")
-selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs)
-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}")
 
 # Tabs
-tab1, tab2, tab3, tab4 = st.tabs(["Build Titan & Camera Snap 🌱📷", "Fine-Tune Titan 🔧", "Test Titan 🧪", "Agentic RAG Party 🌐"])
+tab1, tab2, tab3, tab4 = st.tabs(["Camera Snap 📷", "Test OCR 🔍", "Test Image Gen 🎨", "Test Line Drawings ✏️"])
 
 with tab1:
-    st.header("Build Titan & Camera Snap 🌱📷")
-    st.subheader("Build Titan 🌱")
-    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
-    base_model_options = {
-        "Causal LM": ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"],
-        "Diffusion": [
-            "OFA-Sys/small-stable-diffusion-v0 (LDM/Conditional, ~300 MB)", 
-            "google/ddpm-ema-celebahq-256 (DDPM/SDE/Autoregressive Proxy, ~280 MB)"
-        ]
-    }
-    base_model = st.selectbox("Select Tiny Model", base_model_options[model_type])
-    action = st.radio("Action", ["Use Model", "Download Model"], index=0 if "Causal LM" in model_type else 1)
-    model_name = st.text_input("Model Name (for Download)", f"tiny-titan-{int(time.time())}") if action == "Download Model" else None
-    if st.button(f"{action} ⬇️"):
-        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name or base_model.split(" ")[0], base_model=base_model.split(" ")[0], size="small")
-        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
-        model_type_for_diffusion = "StableDiffusion" if "small-stable-diffusion" in base_model else "DDPM"
-        builder.load_model(base_model.split(" ")[0], config, model_type_for_diffusion, download=action == "Download Model")
-        if action == "Download Model":
-            builder.save_model(config.model_path)
-        st.session_state['builder'] = builder
-        st.session_state['model_loaded'] = True
-        st.rerun()
-
-    st.subheader("Camera Snap 📷")
+    st.header("Camera Snap 📷")
     slice_count = st.number_input("Image Slice Count", min_value=1, max_value=20, value=10)
-    video_length = st.number_input("Video Length (seconds)", min_value=1, max_value=30, value=10)
     cols = st.columns(2)
     with cols[0]:
         st.subheader("Camera 0")
@@ -429,97 +160,52 @@ with tab1:
                 st.image(Image.open(frame), caption=frame, use_container_width=True)
 
 with tab2:
-    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):
-            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")
-                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.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
-        elif isinstance(st.session_state['builder'], DiffusionBuilder):
-            captured_images = get_gallery_files(["png"])
-            if len(captured_images) >= 2:
-                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_images[:min(len(captured_images), 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.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)
+    st.header("Test OCR 🔍")
+    captured_images = get_gallery_files(["png"])
+    if captured_images:
+        selected_image = st.selectbox("Select Image", captured_images)
+        image = Image.open(selected_image)
+        st.image(image, caption="Input Image", use_container_width=True)
+        ocr_model = st.selectbox("Select OCR Model", ["Qwen2-VL-OCR-2B", "GOT-OCR2_0"])
+        prompt = st.text_area("Prompt", "Extract text from the image")
+        if st.button("Run OCR 🚀"):
+            if ocr_model == "Qwen2-VL-OCR-2B":
+                processor, model = load_ocr_qwen2vl()
+                inputs = processor(text=[prompt], images=[image], return_tensors="pt").to("cpu")
+                outputs = model.generate(**inputs, max_new_tokens=1024)
+                text = processor.decode(outputs[0], skip_special_tokens=True)
+            else:  # GOT-OCR2_0
+                tokenizer, model = load_ocr_got()
+                with open(selected_image, "rb") as f:
+                    img_bytes = f.read()
+                img = Image.open(BytesIO(img_bytes))
+                text = model.chat(tokenizer, img, ocr_type='ocr')
+            st.text_area("OCR Result", text, height=200)
 
 with tab3:
-    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:
-        captured_images = get_gallery_files(["png"])
-        if captured_images:
-            selected_image = st.selectbox("Select Image", captured_images)
-            prompt = st.text_area("Enter Text Prompt", f"Superhero {os.path.basename(selected_image).split('.')[0]}")
-            pipeline_options = ["Stable Diffusion (LDM/Conditional)", "DDPM (DDPM/SDE/Autoregressive Proxy)"] if isinstance(st.session_state['builder'], DiffusionBuilder) else ["Causal LM"]
-            selected_pipeline = st.selectbox("Select Pipeline", pipeline_options)
-            if st.button("Run Test 🚀"):
-                if isinstance(st.session_state['builder'], ModelBuilder):
-                    result = st.session_state['builder'].evaluate(prompt)
-                    st.write(f"**Generated Response**: {result}")
-                elif isinstance(st.session_state['builder'], DiffusionBuilder):
-                    if selected_pipeline == "Stable Diffusion (LDM/Conditional)":
-                        image = st.session_state['builder'].generate(prompt)
-                    else:  # DDPM
-                        image = st.session_state['builder'].generate(prompt)
-                    st.image(image, caption=f"Generated from {selected_pipeline}")
+    st.header("Test Image Gen 🎨")
+    captured_images = get_gallery_files(["png"])
+    if captured_images:
+        selected_image = st.selectbox("Select Image", captured_images)
+        image = Image.open(selected_image)
+        st.image(image, caption="Reference Image", use_container_width=True)
+        prompt = st.text_area("Prompt", "Generate a similar superhero image")
+        if st.button("Run Image Gen 🚀"):
+            pipeline = load_image_gen()
+            gen_image = pipeline(prompt, num_inference_steps=50).images[0]
+            st.image(gen_image, caption="Generated Image", use_container_width=True)
 
 with tab4:
-    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.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.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']}")
-
-# Display Logs
-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.header("Test Line Drawings ✏️")
+    captured_images = get_gallery_files(["png"])
+    if captured_images:
+        selected_image = st.selectbox("Select Image", captured_images)
+        image = Image.open(selected_image)
+        st.image(image, caption="Input Image", use_container_width=True)
+        if st.button("Run Line Drawing 🚀"):
+            edge_fn = load_line_drawer()
+            line_drawing = edge_fn(image)
+            st.image(line_drawing, caption="Line Drawing", use_container_width=True)
 
 # Initial Gallery Update
 update_gallery()