bak.05022025.app.py

# --- Combined Imports ------------------------------------
import io
import os
import re
import base64
import glob
import logging
import random
import shutil
import time
import zipfile
import json
import asyncio
import aiofiles

from datetime import datetime
from collections import Counter
from dataclasses import dataclass, field
from io import BytesIO
from typing import Optional, List, Dict, Any

import pandas as pd
import pytz
import streamlit as st
from PIL import Image, ImageDraw # Added ImageDraw
from reportlab.pdfgen import canvas
from reportlab.lib.utils import ImageReader
from reportlab.lib.pagesizes import letter # Default page size
import fitz # PyMuPDF

# --- Hugging Face Imports ---
from huggingface_hub import InferenceClient, HfApi, list_models
from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Import specific exceptions



# --- App Configuration -----------------------------------
st.set_page_config(
    page_title="Vision & Layout Titans (HF) 🚀🖼️",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://huggingface.co/docs',
        'Report a Bug': None, # Replace with your bug report link if desired
        'About': "Combined App: Image->PDF Layout + Hugging Face Powered AI Tools 🌌"
    }
)


# Conditional imports for optional/heavy libraries
try:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline
    # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.)
    _transformers_available = True
except ImportError:
    _transformers_available = False
    st.sidebar.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.")

try:
    from diffusers import StableDiffusionPipeline
    _diffusers_available = True
except ImportError:
    _diffusers_available = False
    # Don't show warning if transformers also missing, handled above
    if _transformers_available:
        st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.")


import requests # Keep requests import

# --- Logging Setup ---------------------------------------
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())

# --- Environment Variables & Constants -------------------
HF_TOKEN = os.getenv("HF_TOKEN")
DEFAULT_PROVIDER = "hf-inference"
# Model List (curated, similar to Gradio example) - can be updated
FEATURED_MODELS_LIST = [
    "meta-llama/Meta-Llama-3.1-8B-Instruct", # Updated Llama model
    "mistralai/Mistral-7B-Instruct-v0.3",
    "google/gemma-2-9b-it", # Added Gemma 2
    "Qwen/Qwen2-7B-Instruct", # Added Qwen2
    "microsoft/Phi-3-mini-4k-instruct",
    "HuggingFaceH4/zephyr-7b-beta",
    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", # Larger Mixture of Experts
    # Add a smaller option
    "HuggingFaceTB/SmolLM-1.7B-Instruct"
]
# Add common vision models if planning local loading
VISION_MODELS_LIST = [
    "Salesforce/blip-image-captioning-large",
    "microsoft/trocr-large-handwritten", # OCR model
    "llava-hf/llava-1.5-7b-hf", # Vision Language Model
    "google/vit-base-patch16-224", # Basic Vision Transformer
]
DIFFUSION_MODELS_LIST = [
    "stabilityai/stable-diffusion-xl-base-1.0", # Common SDXL
    "runwayml/stable-diffusion-v1-5", # Classic SD 1.5
    "OFA-Sys/small-stable-diffusion-v0", # Tiny diffusion
]


# --- Session State Initialization (Combined & Updated) ---
# Layout PDF specific
st.session_state.setdefault('layout_snapshots', [])
st.session_state.setdefault('layout_new_uploads', [])

# General App State
st.session_state.setdefault('history', [])
st.session_state.setdefault('processing', {})
st.session_state.setdefault('asset_checkboxes', {})
st.session_state.setdefault('downloaded_pdfs', {})
st.session_state.setdefault('unique_counter', 0)
st.session_state.setdefault('cam0_file', None)
st.session_state.setdefault('cam1_file', None)
st.session_state.setdefault('characters', [])
st.session_state.setdefault('char_form_reset_key', 0) # For character form reset
st.session_state.setdefault('gallery_size', 10)

# --- Hugging Face & Local Model State ---
st.session_state.setdefault('hf_inference_client', None) # Store initialized client
st.session_state.setdefault('hf_provider', DEFAULT_PROVIDER)
st.session_state.setdefault('hf_custom_key', "")
st.session_state.setdefault('hf_selected_api_model', FEATURED_MODELS_LIST[0]) # Default API model
st.session_state.setdefault('hf_custom_api_model', "") # User override for API model

# Local Model Management
st.session_state.setdefault('local_models', {}) # Dict to store loaded models: {'path': {'model': obj, 'tokenizer': obj, 'type': 'causal/vision/etc'}}
st.session_state.setdefault('selected_local_model_path', None) # Path of the currently active local model

# Inference Parameters (shared for API and local where applicable)
st.session_state.setdefault('gen_max_tokens', 512)
st.session_state.setdefault('gen_temperature', 0.7)
st.session_state.setdefault('gen_top_p', 0.95)
st.session_state.setdefault('gen_frequency_penalty', 0.0)
st.session_state.setdefault('gen_seed', -1) # -1 for random

if 'asset_gallery_container' not in st.session_state:
    st.session_state['asset_gallery_container'] = st.sidebar.empty()

# --- Dataclasses (Refined for Local Models) -------------
@dataclass
class LocalModelConfig:
    name: str                   # User-defined local name
    hf_id: str                  # Hugging Face model ID used for download
    model_type: str             # 'causal', 'vision', 'diffusion', 'ocr', etc.
    size_category: str = "unknown" # e.g., 'small', 'medium', 'large'
    domain: Optional[str] = None
    local_path: str = field(init=False) # Path where it's saved

    def __post_init__(self):
        # Define local path based on type and name
        type_folder = f"{self.model_type}_models"
        safe_name = re.sub(r'[^\w\-]+', '_', self.name) # Sanitize name for path
        self.local_path = os.path.join(type_folder, safe_name)

    def get_full_path(self):
        return os.path.abspath(self.local_path)

# (Keep DiffusionConfig if still using diffusers library separately)
@dataclass
class DiffusionConfig: # Kept for clarity in diffusion tab if needed
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"


# --- Helper Functions (Combined and refined) -------------
# (Keep generate_filename, pdf_url_to_filename, get_download_link, zip_directory)
# ... (previous helper functions like generate_filename, pdf_url_to_filename etc. are assumed here) ...
def generate_filename(sequence, ext="png"):
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    safe_sequence = re.sub(r'[^\w\-]+', '_', str(sequence))
    return f"{safe_sequence}_{timestamp}.{ext}"

def pdf_url_to_filename(url):
    name = re.sub(r'^https?://', '', url)
    name = re.sub(r'[<>:"/\\|?*]', '_', name)
    return name[:100] + ".pdf" # Limit length

def get_download_link(file_path, mime_type="application/octet-stream", label="Download"):
    if not os.path.exists(file_path): return f"{label} (File not found)"
    try:
        with open(file_path, "rb") as f: file_bytes = f.read()
        b64 = base64.b64encode(file_bytes).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"Error creating download link for {file_path}: {e}")
        return f"{label} (Error)"

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:
                file_path = os.path.join(root, file)
                zipf.write(file_path, os.path.relpath(file_path, os.path.dirname(directory_path)))

def get_local_model_paths(model_type="causal"):
    """Gets paths of locally saved models of a specific type."""
    pattern = f"{model_type}_models/*"
    dirs = [d for d in glob.glob(pattern) if os.path.isdir(d)]
    return dirs

def get_gallery_files(file_types=("png", "pdf", "jpg", "jpeg", "md", "txt")):
    all_files = set()
    for ext in file_types:
        all_files.update(glob.glob(f"*.{ext.lower()}"))
        all_files.update(glob.glob(f"*.{ext.upper()}"))
    return sorted(list(all_files))

def get_pdf_files():
    return sorted(glob.glob("*.pdf") + glob.glob("*.PDF"))

def download_pdf(url, output_path):
    try:
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url, stream=True, timeout=20, headers=headers)
        response.raise_for_status()
        with open(output_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
        logger.info(f"Successfully downloaded {url} to {output_path}")
        return True
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to download {url}: {e}")
        if os.path.exists(output_path): 
            try: 
                os.remove(output_path) 
            except: 
                pass
        return False
    except Exception as e:
        logger.error(f"An unexpected error occurred during download of {url}: {e}")
        if os.path.exists(output_path): 
            try: 
                os.remove(output_path) 
            except: 
                pass
        return False

# (Keep process_pdf_snapshot - it doesn't use AI)
async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0):
    start_time = time.time()
    status_placeholder = st.empty()
    status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... (0s)")
    output_files = []
    try:
        doc = fitz.open(pdf_path)
        matrix = fitz.Matrix(resolution_factor, resolution_factor)
        num_pages_to_process = 0
        if mode == "single": num_pages_to_process = min(1, len(doc))
        elif mode == "twopage": num_pages_to_process = min(2, len(doc))
        elif mode == "allpages": num_pages_to_process = len(doc)

        for i in range(num_pages_to_process):
            page_start_time = time.time()
            page = doc[i]
            pix = page.get_pixmap(matrix=matrix)
            base_name = os.path.splitext(os.path.basename(pdf_path))[0]
            output_file = generate_filename(f"{base_name}_pg{i+1}_{mode}", "png")
            await asyncio.to_thread(pix.save, output_file)
            output_files.append(output_file)
            elapsed_page = int(time.time() - page_start_time)
            status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... Page {i+1}/{num_pages_to_process} done ({elapsed_page}s)")
            await asyncio.sleep(0.01)

        doc.close()
        elapsed = int(time.time() - start_time)
        status_placeholder.success(f"PDF Snapshot ({mode}, {len(output_files)} files) completed in {elapsed}s!")
        return output_files
    except Exception as e:
        logger.error(f"Failed to process PDF snapshot for {pdf_path}: {e}")
        status_placeholder.error(f"Failed to process PDF {os.path.basename(pdf_path)}: {e}")
        for f in output_files:
            if os.path.exists(f): os.remove(f)
        return []


# --- HF Inference Client Management ---
def get_hf_client() -> Optional[InferenceClient]:
    """Gets or initializes the Hugging Face Inference Client based on session state."""
    provider = st.session_state.hf_provider
    custom_key = st.session_state.hf_custom_key.strip()
    token_to_use = custom_key if custom_key else HF_TOKEN

    if not token_to_use and provider != "hf-inference":
        st.error(f"Provider '{provider}' requires a Hugging Face API token (either via HF_TOKEN env var or custom key).")
        return None
    if provider == "hf-inference" and not token_to_use:
         logger.warning("Using hf-inference provider without a token. Rate limits may apply.")
         token_to_use = None # Explicitly set to None for public inference API

    # Check if client needs re-initialization
    current_client = st.session_state.get('hf_inference_client')
    # Simple check: re-init if provider or token presence changes
    needs_reinit = True
    if current_client:
         # Basic check, more robust checks could compare client._token etc. if needed
         # This assumes provider and token status are the key determinants
         client_uses_custom = hasattr(current_client, '_token') and current_client._token == custom_key
         client_uses_default = hasattr(current_client, '_token') and current_client._token == HF_TOKEN
         client_uses_no_token = not hasattr(current_client, '_token') or current_client._token is None

         if current_client.provider == provider:
             if custom_key and client_uses_custom: needs_reinit = False
             elif not custom_key and HF_TOKEN and client_uses_default: needs_reinit = False
             elif not custom_key and not HF_TOKEN and client_uses_no_token: needs_reinit = False


    if needs_reinit:
        try:
            logger.info(f"Initializing InferenceClient for provider: {provider}. Token source: {'Custom Key' if custom_key else ('HF_TOKEN' if HF_TOKEN else 'None')}")
            st.session_state.hf_inference_client = InferenceClient(token=token_to_use, provider=provider)
            logger.info("InferenceClient initialized successfully.")
        except Exception as e:
            st.error(f"Failed to initialize Hugging Face client for provider {provider}: {e}")
            logger.error(f"InferenceClient initialization failed: {e}")
            st.session_state.hf_inference_client = None

    return st.session_state.hf_inference_client

# --- HF/Local Model Processing Functions (Replaced OpenAI ones) ---

def process_text_hf(text: str, prompt: str, use_api: bool) -> str:
    """Processes text using either HF Inference API or a loaded local model."""
    status_placeholder = st.empty()
    start_time = time.time()
    result_text = ""

    # --- Prepare Parameters ---
    params = {
        "max_new_tokens": st.session_state.gen_max_tokens, # Note: HF uses max_new_tokens typically
        "temperature": st.session_state.gen_temperature,
        "top_p": st.session_state.gen_top_p,
        "repetition_penalty": st.session_state.gen_frequency_penalty + 1.0, # Adjust HF param name if needed
    }
    seed = st.session_state.gen_seed
    if seed != -1: params["seed"] = seed

    # --- Prepare Messages ---
    # Simple system prompt + user prompt structure
    # More complex chat history could be added here if needed
    system_prompt = "You are a helpful assistant. Process the following text based on the user's request." # Default, consider making configurable
    full_prompt = f"{prompt}\n\n---\n\n{text}"
    # Basic message format for many models, adjust if needed per model type
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": full_prompt}
    ]


    if use_api:
        # --- Use Hugging Face Inference API ---
        status_placeholder.info("Processing text using Hugging Face API...")
        client = get_hf_client()
        if not client:
            return "Error: Hugging Face client not available or configured correctly."

        model_id = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
        if not model_id:
            return "Error: No Hugging Face API model selected or specified."
        status_placeholder.info(f"Using API Model: {model_id}")

        try:
            # Non-streaming for simplicity in Streamlit integration first
            response = client.chat_completion(
                model=model_id,
                messages=messages,
                max_tokens=params['max_new_tokens'], # chat_completion uses max_tokens
                temperature=params['temperature'],
                top_p=params['top_p'],
                # Add other params if supported by client.chat_completion
            )
            result_text = response.choices[0].message.content or ""
            logger.info(f"HF API text processing successful for model {model_id}.")

        except Exception as e:
            logger.error(f"HF API text processing failed for model {model_id}: {e}")
            result_text = f"Error during Hugging Face API inference: {str(e)}"

    else:
        # --- Use Loaded Local Model ---
        status_placeholder.info("Processing text using local model...")
        if not _transformers_available:
            return "Error: Transformers library not available for local models."

        model_path = st.session_state.get('selected_local_model_path')
        if not model_path or model_path not in st.session_state.get('local_models', {}):
            return "Error: No suitable local model selected or loaded."

        local_model_data = st.session_state['local_models'][model_path]
        if local_model_data.get('type') != 'causal':
             return f"Error: Loaded model '{os.path.basename(model_path)}' is not a Causal LM."

        status_placeholder.info(f"Using Local Model: {os.path.basename(model_path)}")
        model = local_model_data.get('model')
        tokenizer = local_model_data.get('tokenizer')

        if not model or not tokenizer:
             return f"Error: Model or tokenizer not found for {os.path.basename(model_path)}."

        try:
            # Prepare input for local transformers model
            # Handle chat template if available, otherwise basic concatenation
            try:
                prompt_for_model = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            except Exception: # Fallback if template fails or doesn't exist
                 logger.warning(f"Could not apply chat template for {model_path}. Using basic formatting.")
                 prompt_for_model = f"System: {system_prompt}\nUser: {full_prompt}\nAssistant:"

            inputs = tokenizer(prompt_for_model, return_tensors="pt", padding=True, truncation=True, max_length=params['max_new_tokens'] * 2) # Heuristic length limit
            # Move inputs to the same device as the model
            inputs = {k: v.to(model.device) for k, v in inputs.items()}

            # Generate
            # Ensure generate parameters match transformers' expected names
            generate_params = {
                "max_new_tokens": params['max_new_tokens'],
                "temperature": params['temperature'],
                "top_p": params['top_p'],
                "repetition_penalty": params.get('repetition_penalty', 1.0), # Use adjusted name
                "do_sample": True if params['temperature'] > 0.1 else False, # Required for temp/top_p
                "pad_token_id": tokenizer.eos_token_id # Avoid PAD warning
            }
            if 'seed' in params: pass # Seed handling can be complex with transformers, often set globally

            with torch.no_grad(): # Disable gradient calculation for inference
                 outputs = model.generate(**inputs, **generate_params)

            # Decode the output, skipping special tokens and the prompt
            # output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            # More robust decoding: only decode the newly generated part
            input_length = inputs['input_ids'].shape[1]
            generated_ids = outputs[0][input_length:]
            result_text = tokenizer.decode(generated_ids, skip_special_tokens=True)

            logger.info(f"Local text processing successful for model {model_path}.")

        except Exception as e:
            logger.error(f"Local text processing failed for model {model_path}: {e}")
            result_text = f"Error during local model inference: {str(e)}"


    elapsed = int(time.time() - start_time)
    status_placeholder.success(f"Text processing completed in {elapsed}s.")
    return result_text


# --- Image Processing (Placeholder/Basic Implementation) ---
# This needs significant work depending on the chosen vision model type
def process_image_hf(image: Image.Image, prompt: str, use_api: bool) -> str:
    """Processes an image using either HF Inference API or a local model."""
    status_placeholder = st.empty()
    start_time = time.time()
    result_text = "[Image processing not fully implemented with HF models yet]"

    if use_api:
        # --- Use HF API (Basic Image-to-Text Example) ---
        status_placeholder.info("Processing image using Hugging Face API (Image-to-Text)...")
        client = get_hf_client()
        if not client: return "Error: HF client not configured."

        # Convert PIL image to bytes
        buffered = BytesIO()
        image.save(buffered, format="PNG" if image.format != 'JPEG' else 'JPEG')
        img_bytes = buffered.getvalue()

        try:
            # Example using a generic image-to-text model via API
            # NOTE: This does NOT use the 'prompt' effectively like VQA models.
            # Need to select an appropriate model ID known for image captioning.
            # Using a default BLIP model for demonstration.
            captioning_model_id = "Salesforce/blip-image-captioning-large"
            status_placeholder.info(f"Using API Image-to-Text Model: {captioning_model_id}")

            response_list = client.image_to_text(data=img_bytes, model=captioning_model_id)

            if response_list and isinstance(response_list, list) and 'generated_text' in response_list[0]:
                result_text = f"API Caption ({captioning_model_id}): {response_list[0]['generated_text']}\n\n(Note: API call did not use custom prompt: '{prompt}')"
                logger.info(f"HF API image captioning successful for model {captioning_model_id}.")
            else:
                 result_text = "Error: Unexpected response format from image-to-text API."
                 logger.warning(f"Unexpected API response for image-to-text: {response_list}")

        except Exception as e:
            logger.error(f"HF API image processing failed: {e}")
            result_text = f"Error during Hugging Face API image inference: {str(e)}"

    else:
        # --- Use Local Vision Model ---
        status_placeholder.info("Processing image using local model...")
        if not _transformers_available: return "Error: Transformers library needed."

        model_path = st.session_state.get('selected_local_model_path')
        if not model_path or model_path not in st.session_state.get('local_models', {}):
            return "Error: No suitable local model selected or loaded."

        local_model_data = st.session_state['local_models'][model_path]
        model_type = local_model_data.get('type')

        # --- Placeholder Logic - Requires Specific Model Implementation ---
        if model_type == 'vision': # General VQA or Captioning
             status_placeholder.warning(f"Local Vision Model ({os.path.basename(model_path)}): Processing logic depends heavily on the specific model architecture (e.g., LLaVA, BLIP). Placeholder implementation.")
             # Example: Needs processor + model.generate based on model type
             # processor = local_model_data.get('processor')
             # model = local_model_data.get('model')
             # if processor and model:
             #     try:
             #         # inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
             #         # generated_ids = model.generate(**inputs, max_new_tokens=...)
             #         # result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
             #         result_text = f"[Local vision processing for {os.path.basename(model_path)} needs specific implementation based on its type.] Prompt was: {prompt}"
             #     except Exception as e:
             #         result_text = f"Error during local vision model inference: {e}"
             # else:
             #      result_text = "Error: Processor or model missing for local vision task."
             result_text = f"[Local vision processing for {os.path.basename(model_path)} needs specific implementation based on its type.] Prompt was: {prompt}" # Placeholder

        elif model_type == 'ocr': # OCR Specific Model
            status_placeholder.warning(f"Local OCR Model ({os.path.basename(model_path)}): Placeholder implementation.")
            # Example for TrOCR style models
            # processor = local_model_data.get('processor')
            # model = local_model_data.get('model')
            # if processor and model:
            #      try:
            #         # pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(model.device)
            #         # generated_ids = model.generate(pixel_values, max_new_tokens=...)
            #         # result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
            #         result_text = f"[Local OCR processing for {os.path.basename(model_path)} needs specific implementation.]"
            #      except Exception as e:
            #         result_text = f"Error during local OCR model inference: {e}"
            # else:
            #      result_text = "Error: Processor or model missing for local OCR task."
            result_text = f"[Local OCR processing for {os.path.basename(model_path)} needs specific implementation.]" # Placeholder
        else:
             result_text = f"Error: Loaded model '{os.path.basename(model_path)}' is not a recognized vision/OCR type for this function."


    elapsed = int(time.time() - start_time)
    status_placeholder.success(f"Image processing attempt completed in {elapsed}s.")
    return result_text

# Basic OCR function using the image processor above
async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool) -> str:
    """ Performs OCR using the process_image_hf function framework. """
    # Simple prompt for OCR task
    ocr_prompt = "Extract text content from this image."
    result = process_image_hf(image, ocr_prompt, use_api)

    # Save the result if it looks like text (basic check)
    if result and not result.startswith("Error") and not result.startswith("["):
        try:
            async with aiofiles.open(output_file, "w", encoding='utf-8') as f:
                await f.write(result)
            logger.info(f"HF OCR result saved to {output_file}")
        except IOError as e:
             logger.error(f"Failed to save HF OCR output to {output_file}: {e}")
             result += f"\n[Error saving file: {e}]" # Append error to result if save fails
    elif os.path.exists(output_file):
         # Remove file if processing failed or was just a placeholder message
         try: os.remove(output_file)
         except OSError: pass

    return result


# --- Character Functions (Keep from previous) -----------
# ... (randomize_character_content, save_character, load_characters are assumed here) ...
def randomize_character_content():
    intro_templates = [
        "{char} is a valiant knight...", "{char} is a mischievous thief...",
        "{char} is a wise scholar...", "{char} is a fiery warrior...", "{char} is a gentle healer..."
    ]
    greeting_templates = [
        "'I am from the knight's guild...'", "'I heard you needed help—name’s {char}...",
        "'Oh, hello! I’m {char}, didn’t see you there...'", "'I’m {char}, and I’m here to fight...'",
        "'I’m {char}, here to heal...'" ]
    name = f"Character_{random.randint(1000, 9999)}"
    gender = random.choice(["Male", "Female"])
    intro = random.choice(intro_templates).format(char=name)
    greeting = random.choice(greeting_templates).format(char=name)
    return name, gender, intro, greeting

def save_character(character_data):
    characters = st.session_state.get('characters', [])
    if any(c['name'] == character_data['name'] for c in characters):
         st.error(f"Character name '{character_data['name']}' already exists.")
         return False
    characters.append(character_data)
    st.session_state['characters'] = characters
    try:
        with open("characters.json", "w", encoding='utf-8') as f: json.dump(characters, f, indent=2)
        logger.info(f"Saved character: {character_data['name']}")
        return True
    except IOError as e:
        logger.error(f"Failed to save characters.json: {e}")
        st.error(f"Failed to save character file: {e}")
        return False

def load_characters():
    if not os.path.exists("characters.json"): st.session_state['characters'] = []; return
    try:
        with open("characters.json", "r", encoding='utf-8') as f: characters = json.load(f)
        if isinstance(characters, list): st.session_state['characters'] = characters; logger.info(f"Loaded {len(characters)} characters.")
        else: st.session_state['characters'] = []; logger.warning("characters.json is not a list, resetting."); os.remove("characters.json")
    except (json.JSONDecodeError, IOError) as e:
        logger.error(f"Failed to load or decode characters.json: {e}")
        st.error(f"Error loading character file: {e}. Starting fresh.")
        st.session_state['characters'] = []
        try:
            corrupt_filename = f"characters_corrupt_{int(time.time())}.json"
            shutil.copy("characters.json", corrupt_filename); logger.info(f"Backed up corrupted character file to {corrupt_filename}"); os.remove("characters.json")
        except Exception as backup_e: logger.error(f"Could not backup corrupted character file: {backup_e}")


# --- Utility: Clean stems (Keep from previous) ----------
def clean_stem(fn: str) -> str:
    name = os.path.splitext(os.path.basename(fn))[0]
    name = name.replace('-', ' ').replace('_', ' ')
    return name.strip().title()


# --- PDF Creation: Image Sized + Captions (Keep from previous) ---
def make_image_sized_pdf(sources):
    if not sources: st.warning("No image sources provided for PDF generation."); return None
    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=letter) # Default letter
    try:
        for idx, src in enumerate(sources, start=1):
            status_placeholder = st.empty()
            status_placeholder.info(f"Adding page {idx}/{len(sources)}: {os.path.basename(str(src))}...")
            try:
                filename = f'page_{idx}'
                if isinstance(src, str):
                    if not os.path.exists(src): logger.warning(f"Image file not found: {src}. Skipping."); status_placeholder.warning(f"Skipping missing file: {os.path.basename(src)}"); continue
                    img_obj = Image.open(src); filename = os.path.basename(src)
                else:
                    src.seek(0); img_obj = Image.open(src); filename = getattr(src, 'name', f'uploaded_image_{idx}'); src.seek(0)

                with img_obj:
                    iw, ih = img_obj.size
                    if iw <= 0 or ih <= 0: logger.warning(f"Invalid image dimensions ({iw}x{ih}) for {filename}. Skipping."); status_placeholder.warning(f"Skipping invalid image: {filename}"); continue
                    cap_h = 30; pw, ph = iw, ih + cap_h
                    c.setPageSize((pw, ph))
                    img_reader = ImageReader(img_obj)
                    c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto')
                    caption = clean_stem(filename); c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, caption)
                    c.setFont('Helvetica', 8); c.setFillColorRGB(0.5, 0.5, 0.5); c.drawRightString(pw - 10, 8, f"Page {idx}")
                    c.showPage()
                    status_placeholder.success(f"Added page {idx}/{len(sources)}: {filename}")

            except (IOError, OSError, UnidentifiedImageError) as img_err: logger.error(f"Error processing image {src}: {img_err}"); status_placeholder.error(f"Error adding page {idx}: {img_err}")
            except Exception as e: logger.error(f"Unexpected error adding page {idx} ({src}): {e}"); status_placeholder.error(f"Unexpected error on page {idx}: {e}")

        c.save(); buf.seek(0)
        if buf.getbuffer().nbytes < 100: st.error("PDF generation resulted in an empty file."); return None
        return buf.getvalue()
    except Exception as e:
        logger.error(f"Fatal error during PDF generation: {e}")
        st.error(f"PDF Generation Failed: {e}")
        return None


# --- Sidebar Gallery Update Function (MODIFIED) --------
def update_gallery():
    st.sidebar.markdown("### Asset Gallery 📸📖")

    all_files = get_gallery_files() # Get currently available files

    if not all_files:
        st.sidebar.info("No assets (images, PDFs, text files) found yet.")
        return

    st.sidebar.caption(f"Found {len(all_files)} assets:")

    for idx, file in enumerate(all_files):
        st.session_state['unique_counter'] += 1
        unique_id = st.session_state['unique_counter']
        item_key_base = f"gallery_item_{os.path.basename(file)}_{unique_id}"
        basename = os.path.basename(file)
        st.sidebar.markdown(f"**{basename}**") # Display filename clearly

        try:
            file_ext = os.path.splitext(file)[1].lower()
            # Display previews
            if file_ext in ['.png', '.jpg', '.jpeg']:
                 # Add expander for large galleries
                 with st.sidebar.expander("Preview", expanded=False):
                      st.image(Image.open(file), use_container_width=True)
            elif file_ext == '.pdf':
                 with st.sidebar.expander("Preview (Page 1)", expanded=False):
                      doc = fitz.open(file)
                      if len(doc) > 0:
                          pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) # Smaller preview
                          img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                          st.image(img, use_container_width=True)
                      else:
                          st.warning("Empty PDF")
                      doc.close()
            elif file_ext in ['.md', '.txt']:
                 with st.sidebar.expander("Preview (Start)", expanded=False):
                      with open(file, 'r', encoding='utf-8', errors='ignore') as f:
                          content_preview = f.read(200) # Show first 200 chars
                      st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text')

            # --- Actions for the file (Select, Download, Delete) ---
            action_cols = st.sidebar.columns(3) # Use columns for buttons
            with action_cols[0]:
                 checkbox_key = f"cb_{item_key_base}"
                 st.session_state['asset_checkboxes'][file] = st.checkbox(
                     "Select",
                     value=st.session_state['asset_checkboxes'].get(file, False),
                     key=checkbox_key
                 )
            with action_cols[1]:
                mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'}
                mime_type = mime_map.get(file_ext, "application/octet-stream")
                # Use button for download to avoid complex HTML link generation issues sometimes
                dl_key = f"dl_{item_key_base}"
                try:
                    with open(file, "rb") as fp:
                        st.download_button(
                             label="📥",
                             data=fp,
                             file_name=basename,
                             mime=mime_type,
                             key=dl_key,
                             help="Download this file"
                         )
                except Exception as dl_e:
                     st.error(f"DL Err: {dl_e}")

            with action_cols[2]:
                delete_key = f"del_{item_key_base}"
                if st.button("🗑️", key=delete_key, help=f"Delete {basename}"):
                    try:
                        os.remove(file)
                        st.session_state['asset_checkboxes'].pop(file, None) # Remove from selection state
                        # Remove from layout_snapshots if present
                        if file in st.session_state.get('layout_snapshots', []):
                            st.session_state['layout_snapshots'].remove(file)
                        logger.info(f"Deleted asset: {file}")
                        st.toast(f"Deleted {basename}!", icon="✅") # Use toast for less intrusive feedback
                        # REMOVED st.rerun() - Rely on file watcher
                    except OSError as e:
                        logger.error(f"Error deleting file {file}: {e}")
                        st.error(f"Could not delete {basename}")
                    # Trigger a rerun MANUALLY after deletion completes if file watcher is unreliable
                    st.rerun()


        except (fitz.fitz.FileNotFoundError, FileNotFoundError):
             st.sidebar.error(f"File not found: {basename}")
             st.session_state['asset_checkboxes'].pop(file, None) # Clean up state
        except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err:
             st.sidebar.error(f"Corrupt PDF: {basename}")
             logger.warning(f"Error opening PDF {file}: {pdf_err}")
        except UnidentifiedImageError:
            st.sidebar.error(f"Invalid Image: {basename}")
            logger.warning(f"Cannot identify image file {file}")
        except Exception as e:
            st.sidebar.error(f"Error: {basename}")
            logger.error(f"Error displaying asset {file}: {e}")

        st.sidebar.markdown("---") # Separator between items

# --- UI Elements -----------------------------------------

# --- Sidebar: HF Inference Settings ---
st.sidebar.subheader("🤖 Hugging Face Settings")
st.sidebar.markdown("Configure API inference or select local models.")

# API Settings Expander
with st.sidebar.expander("API Inference Settings", expanded=False):
    st.session_state.hf_custom_key = st.text_input(
        "Custom HF Token (BYOK)",
        value=st.session_state.get('hf_custom_key', ""),
        type="password",
        key="hf_custom_key_input",
        help="Enter your Hugging Face API token. Overrides HF_TOKEN env var."
    )
    token_status = "Custom Key Set" if st.session_state.hf_custom_key else ("Default HF_TOKEN Set" if HF_TOKEN else "No Token Set")
    st.caption(f"Token Status: {token_status}")

    providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
    st.session_state.hf_provider = st.selectbox(
        "Inference Provider",
        options=providers_list,
        index=providers_list.index(st.session_state.get('hf_provider', DEFAULT_PROVIDER)),
        key="hf_provider_select",
        help="Select the backend provider. Some require specific API keys."
    )
    # Validate provider based on key (simple validation)
    if not st.session_state.hf_custom_key and not HF_TOKEN and st.session_state.hf_provider != "hf-inference":
        st.warning(f"Provider '{st.session_state.hf_provider}' may require a token. Using 'hf-inference' may work without a token but with rate limits.")

    # API Model Selection
    st.session_state.hf_custom_api_model = st.text_input(
        "Custom API Model ID",
        value=st.session_state.get('hf_custom_api_model', ""),
        key="hf_custom_model_input",
        placeholder="e.g., google/gemma-2-9b-it",
        help="Overrides the featured model selection below if provided."
    )
    # Use custom if provided, otherwise use the selected featured model
    effective_api_model = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model

    st.session_state.hf_selected_api_model = st.selectbox(
        "Featured API Model",
        options=FEATURED_MODELS_LIST,
        index=FEATURED_MODELS_LIST.index(st.session_state.get('hf_selected_api_model', FEATURED_MODELS_LIST[0])),
        key="hf_featured_model_select",
        help="Select a common model. Ignored if Custom API Model ID is set."
    )
    st.caption(f"Effective API Model: {effective_api_model}")


# Local Model Selection Expander
with st.sidebar.expander("Local Model Selection", expanded=True):
    if not _transformers_available:
        st.warning("Transformers library not found. Cannot load or use local models.")
    else:
        local_model_options = ["None"] + list(st.session_state.get('local_models', {}).keys())
        current_selection = st.session_state.get('selected_local_model_path')
        # Ensure current selection is valid
        if current_selection not in local_model_options:
             current_selection = "None"

        selected_path = st.selectbox(
            "Active Local Model",
            options=local_model_options,
            index=local_model_options.index(current_selection),
            format_func=lambda x: os.path.basename(x) if x != "None" else "None",
            key="local_model_selector",
            help="Select a model loaded via the 'Build Titan' tab to use for processing."
        )
        st.session_state.selected_local_model_path = selected_path if selected_path != "None" else None

        if st.session_state.selected_local_model_path:
             model_info = st.session_state.local_models[st.session_state.selected_local_model_path]
             st.caption(f"Type: {model_info.get('type', 'Unknown')}")
             st.caption(f"Device: {model_info.get('model').device if model_info.get('model') else 'N/A'}")
        else:
             st.caption("No local model selected.")

# Generation Parameters Expander
with st.sidebar.expander("Generation Parameters", expanded=False):
    st.session_state.gen_max_tokens = st.slider("Max New Tokens", 1, 4096, st.session_state.get('gen_max_tokens', 512), step=1, key="param_max_tokens")
    st.session_state.gen_temperature = st.slider("Temperature", 0.01, 2.0, st.session_state.get('gen_temperature', 0.7), step=0.01, key="param_temp")
    st.session_state.gen_top_p = st.slider("Top-P", 0.01, 1.0, st.session_state.get('gen_top_p', 0.95), step=0.01, key="param_top_p")
    # Note: HF often uses repetition_penalty instead of frequency_penalty. We'll use it here.
    st.session_state.gen_frequency_penalty = st.slider("Repetition Penalty", 1.0, 2.0, st.session_state.get('gen_frequency_penalty', 0.0)+1.0, step=0.05, key="param_repetition", help="1.0 means no penalty.")
    st.session_state.gen_seed = st.slider("Seed", -1, 65535, st.session_state.get('gen_seed', -1), step=1, key="param_seed", help="-1 for random.")



st.sidebar.markdown("---") # Separator before gallery settings

# --- ADDED: Gallery Settings Section ---
st.sidebar.subheader("🖼️ Gallery Settings")
st.slider(
    "Max Items Shown",
    min_value=2,
    max_value=50, # Adjust max if needed
    value=st.session_state.get('gallery_size', 10),
    key="gallery_size_slider", # Keep the key, define it ONCE here
    help="Controls the maximum number of assets displayed in the sidebar gallery."
)
st.session_state.gallery_size = st.session_state.gallery_size_slider # Ensure sync
st.sidebar.markdown("---") # Separator after gallery settings




# --- App Title -------------------------------------------
st.title("Vision & Layout Titans (HF) 🚀🖼️📄")
st.markdown("Combined App: Image-to-PDF Layout + Hugging Face Powered AI Tools")

# --- Main Application Tabs -------------------------------
tab_list = [
    "Image->PDF Layout 🖼️➡️📄", # From App 1
    "Camera Snap 📷",
    "Download PDFs 📥",
    "Build Titan (Local Models) 🌱", # Renamed for clarity
    "Text Process (HF) 📝", # New tab for text
    "Image Process (HF) 🖼️", # New tab for image
    "Test OCR (HF) 🔍", # Renamed
    "Character Editor 🧑‍🎨",
    "Character Gallery 🖼️",
    # Original Tabs (potentially redundant or integrated now):
    # "PDF Process 📄", (Integrated into Text/Image process conceptually)
    # "MD Gallery & Process 📚", (Use Text Process tab)
    # "Test Image Gen 🎨", (Separate Diffusion logic)
]
# Filter out redundant tabs if they are fully replaced
# Example: If MD Gallery is fully handled by Text Process, remove it. For now, keep most.
# Let's keep PDF Process and Image Process separate for clarity of input type, but use the new HF functions
tabs_to_create = [
    "Image->PDF Layout 🖼️➡️📄",
    "Camera Snap 📷",
    "Download PDFs 📥",
    "Build Titan (Local Models) 🌱",
    "PDF Process (HF) 📄", # Use HF functions for PDF pages
    "Image Process (HF) 🖼️",# Use HF functions for images
    "Text Process (HF) 📝", # Use HF functions for MD/TXT files
    "Test OCR (HF) 🔍",   # Use HF OCR logic
    "Test Image Gen (Diffusers) 🎨", # Keep diffusion separate
    "Character Editor 🧑‍🎨",
    "Character Gallery 🖼️",
]

tabs = st.tabs(tabs_to_create)

# --- Tab Implementations ---

# --- Tab 1: Image -> PDF Layout (Keep from previous merge) ---
with tabs[0]:
    # ... (Code from previous merge for this tab remains largely the same) ...
    st.header("Image to PDF Layout Generator")
    st.markdown("Upload or scan images, reorder them, and generate a PDF where each page matches the image dimensions and includes a simple caption.")
    col1, col2 = st.columns(2)
    with col1:
        st.subheader("A. Scan or Upload Images")
        layout_cam = st.camera_input("📸 Scan Document for Layout PDF", key="layout_cam")
        if layout_cam:
            now = datetime.now(pytz.timezone("US/Central"))
            scan_name = generate_filename(f"layout_scan_{now.strftime('%a').upper()}", "png")
            try:
                with open(scan_name, "wb") as f: f.write(layout_cam.getvalue())
                st.image(Image.open(scan_name), caption=f"Scanned: {scan_name}", use_container_width=True)
                if scan_name not in st.session_state['layout_snapshots']: st.session_state['layout_snapshots'].append(scan_name)
                st.success(f"Scan saved as {scan_name}")
                update_gallery();  # Add to gallery
            except Exception as e: st.error(f"Failed to save scan: {e}"); logger.error(f"Failed to save camera scan {scan_name}: {e}")

        layout_uploads = st.file_uploader("📂 Upload PNG/JPG Images for Layout PDF", type=["png","jpg","jpeg"], accept_multiple_files=True, key="layout_uploader")
        if layout_uploads: st.session_state['layout_new_uploads'] = layout_uploads # Store for processing below
    with col2:
        st.subheader("B. Review and Reorder")
        layout_records = []
        processed_snapshots = set()
        # Process snapshots
        for idx, path in enumerate(st.session_state.get('layout_snapshots', [])):
             if path not in processed_snapshots and os.path.exists(path):
                try:
                    with Image.open(path) as im: w, h = im.size; ar = round(w / h, 2) if h > 0 else 0; orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait")
                    layout_records.append({"filename": os.path.basename(path), "source": path, "width": w, "height": h, "aspect_ratio": ar, "orientation": orient, "order": idx, "type": "Scan"})
                    processed_snapshots.add(path)
                except Exception as e: logger.warning(f"Could not process snapshot {path}: {e}"); st.warning(f"Skipping invalid snapshot: {os.path.basename(path)}")
        # Process current uploads
        current_uploads = st.session_state.get('layout_new_uploads', [])
        if current_uploads:
             start_idx = len(layout_records)
             for jdx, f_obj in enumerate(current_uploads, start=start_idx):
                 try:
                     f_obj.seek(0)
                     with Image.open(f_obj) as im: w, h = im.size; ar = round(w / h, 2) if h > 0 else 0; orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait")
                     layout_records.append({"filename": f_obj.name, "source": f_obj, "width": w, "height": h, "aspect_ratio": ar, "orientation": orient, "order": jdx, "type": "Upload"})
                     f_obj.seek(0)
                 except Exception as e: logger.warning(f"Could not process uploaded file {f_obj.name}: {e}"); st.warning(f"Skipping invalid upload: {f_obj.name}")

        if not layout_records: st.info("Scan or upload images using the controls on the left.")
        else:
            layout_df = pd.DataFrame(layout_records); dims = st.multiselect("Include orientations:", options=["Landscape","Portrait","Square"], default=["Landscape","Portrait","Square"], key="layout_dims_filter")
            filtered_df = layout_df[layout_df['orientation'].isin(dims)].copy() if dims else layout_df.copy()
            filtered_df['order'] = filtered_df['order'].astype(int); filtered_df = filtered_df.sort_values('order').reset_index(drop=True)
            st.markdown("Edit 'Order' column or drag rows to set PDF page sequence:")
            edited_df = st.data_editor(filtered_df, column_config={"filename": st.column_config.TextColumn("Filename", disabled=True), "source": None, "width": st.column_config.NumberColumn("Width", disabled=True), "height": st.column_config.NumberColumn("Height", disabled=True), "aspect_ratio": st.column_config.NumberColumn("Aspect Ratio", format="%.2f", disabled=True), "orientation": st.column_config.TextColumn("Orientation", disabled=True), "type": st.column_config.TextColumn("Source Type", disabled=True), "order": st.column_config.NumberColumn("Order", min_value=0, step=1, required=True)}, hide_index=True, use_container_width=True, num_rows="dynamic", key="layout_editor")
            ordered_layout_df = edited_df.sort_values('order').reset_index(drop=True)
            ordered_sources_for_pdf = ordered_layout_df['source'].tolist()

            st.subheader("C. Generate & Download PDF")
            if st.button("🖋️ Generate Image-Sized PDF", key="generate_layout_pdf"):
                if not ordered_sources_for_pdf: st.warning("No images selected or available after filtering.")
                else:
                    with st.spinner("Generating PDF..."): pdf_bytes = make_image_sized_pdf(ordered_sources_for_pdf)
                    if pdf_bytes:
                        now = datetime.now(pytz.timezone("US/Central")); prefix = now.strftime("%Y%m%d-%H%M%p")
                        stems = [clean_stem(s) if isinstance(s, str) else clean_stem(getattr(s, 'name', 'upload')) for s in ordered_sources_for_pdf[:4]]
                        basename = " - ".join(stems) or "Layout"; pdf_fname = f"{prefix}_{basename}.pdf"; pdf_fname = re.sub(r'[^\w\- \.]', '_', pdf_fname)
                        st.success(f"✅ PDF ready: **{pdf_fname}**")
                        st.download_button("⬇️ Download PDF", data=pdf_bytes, file_name=pdf_fname, mime="application/pdf", key="download_layout_pdf")
                        st.markdown("#### Preview First Page")
                        try:
                            doc = fitz.open(stream=pdf_bytes, filetype='pdf')
                            if len(doc) > 0: pix = doc[0].get_pixmap(matrix=fitz.Matrix(1.0, 1.0)); preview_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); st.image(preview_img, caption=f"Preview of {pdf_fname} (Page 1)", use_container_width=True)
                            else: st.warning("Generated PDF appears empty.")
                            doc.close()
                        except Exception as preview_err: st.warning(f"Could not generate PDF preview: {preview_err}"); logger.warning(f"PDF preview error for {pdf_fname}: {preview_err}")
                    else: st.error("PDF generation failed. Check logs or image files.")


# --- Tab 2: Camera Snap (Keep from previous merge) ---
with tabs[1]:
    # ... (Code from previous merge for this tab) ...
    st.header("Camera Snap 📷")
    st.subheader("Single Capture (Adds to General Gallery)")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="main_cam0")
        if cam0_img:
            filename = generate_filename("cam0_snap");
            if st.session_state.get('cam0_file') and os.path.exists(st.session_state['cam0_file']): 
                try: 
                    os.remove(st.session_state['cam0_file']) 
                except OSError: 
                    pass
            try:
                with open(filename, "wb") as f: f.write(cam0_img.getvalue())
                st.session_state['cam0_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 0: {filename}"); st.image(Image.open(filename), caption="Camera 0 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 0: {filename}"); st.success(f"Saved {filename}")
                update_gallery();
            except Exception as e: 
                st.error(f"Failed to save Cam 0 snap: {e}"); logger.error(f"Failed to save Cam 0 snap {filename}: {e}")
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="main_cam1")
        if cam1_img:
            filename = generate_filename("cam1_snap")
            if st.session_state.get('cam1_file') and os.path.exists(st.session_state['cam1_file']): 
                try: 
                    os.remove(st.session_state['cam1_file']) 
                except OSError: 
                    pass
            try:
                with open(filename, "wb") as f: f.write(cam1_img.getvalue())
                st.session_state['cam1_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 1: {filename}"); st.image(Image.open(filename), caption="Camera 1 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 1: {filename}"); st.success(f"Saved {filename}")
                update_gallery();
            except Exception as e: st.error(f"Failed to save Cam 1 snap: {e}"); logger.error(f"Failed to save Cam 1 snap {filename}: {e}")


# --- Tab 3: Download PDFs (Keep from previous merge) ---
with tabs[2]:
    # ... (Code from previous merge for this tab) ...
    st.header("Download PDFs 📥")
    st.markdown("Download PDFs from URLs and optionally create image snapshots.")
    if st.button("Load Example arXiv URLs 📚", key="load_examples"):
        example_urls = ["https://arxiv.org/pdf/2308.03892", "https://arxiv.org/pdf/1706.03762", "https://arxiv.org/pdf/2402.17764", "https://www.clickdimensions.com/links/ACCERL/"]
        st.session_state['pdf_urls_input'] = "\n".join(example_urls)
    url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls_input', ""), height=150, key="pdf_urls_textarea")
    if st.button("Robo-Download PDFs 🤖", key="download_pdfs_button"):
        urls = [url.strip() for url in url_input.strip().split("\n") if url.strip()]
        if not urls: st.warning("Please enter at least one URL.")
        else:
            progress_bar = st.progress(0); status_text = st.empty(); total_urls = len(urls); download_count = 0; existing_pdfs = get_pdf_files()
            for idx, url in enumerate(urls):
                output_path = pdf_url_to_filename(url); status_text.text(f"Processing {idx + 1}/{total_urls}: {os.path.basename(output_path)}..."); progress_bar.progress((idx + 1) / total_urls)
                if output_path in existing_pdfs: st.info(f"Already exists: {os.path.basename(output_path)}"); st.session_state['downloaded_pdfs'][url] = output_path; st.session_state['asset_checkboxes'][output_path] = st.session_state['asset_checkboxes'].get(output_path, False)
                else:
                    if download_pdf(url, output_path): st.session_state['downloaded_pdfs'][url] = output_path; logger.info(f"Downloaded PDF from {url} to {output_path}"); st.session_state['history'].append(f"Downloaded PDF: {output_path}"); st.session_state['asset_checkboxes'][output_path] = False; download_count += 1; existing_pdfs.append(output_path)
                    else: st.error(f"Failed to download: {url}")
            status_text.success(f"Download process complete! Successfully downloaded {download_count} new PDFs.")
            if download_count > 0: update_gallery();

    st.subheader("Create Snapshots from Gallery PDFs")
    snapshot_mode = st.selectbox("Snapshot Mode", ["First Page (High-Res)", "First Two Pages (High-Res)", "All Pages (High-Res)", "First Page (Low-Res Preview)"], key="pdf_snapshot_mode")
    resolution_map = {"First Page (High-Res)": 2.0, "First Two Pages (High-Res)": 2.0, "All Pages (High-Res)": 2.0, "First Page (Low-Res Preview)": 1.0}
    mode_key_map = {"First Page (High-Res)": "single", "First Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages", "First Page (Low-Res Preview)": "single"}
    resolution = resolution_map[snapshot_mode]; mode_key = mode_key_map[snapshot_mode]
    if st.button("Snapshot Selected PDFs 📸", key="snapshot_selected_pdfs"):
        selected_pdfs = [path for path in get_gallery_files(['pdf']) if st.session_state['asset_checkboxes'].get(path, False)]
        if not selected_pdfs: st.warning("No PDFs selected in the sidebar gallery!")
        else:
            st.info(f"Starting snapshot process for {len(selected_pdfs)} selected PDF(s)..."); snapshot_count = 0; total_snapshots_generated = 0
            for pdf_path in selected_pdfs:
                if not os.path.exists(pdf_path): st.warning(f"File not found: {pdf_path}. Skipping."); continue
                new_snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key, resolution))
                if new_snapshots:
                    snapshot_count += 1; total_snapshots_generated += len(new_snapshots)
                    st.write(f"Snapshots for {os.path.basename(pdf_path)}:"); cols = st.columns(3)
                    for i, snap_path in enumerate(new_snapshots):
                         with cols[i % 3]: st.image(Image.open(snap_path), caption=os.path.basename(snap_path), use_container_width=True); st.session_state['asset_checkboxes'][snap_path] = False # Add to gallery
            if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery();
            else: st.warning("No snapshots were generated. Check logs or PDF files.")

# --- Tab 4: Build Titan (Local Models) ---
with tabs[3]:
    st.header("Build Titan (Local Models) 🌱")
    st.markdown("Download and save models from Hugging Face Hub for local use.")

    if not _transformers_available:
        st.error("Transformers library not available. Cannot download or load local models.")
    else:
        build_model_type = st.selectbox(
            "Select Model Type",
            ["Causal LM", "Vision/Multimodal", "OCR", "Diffusion"], # Added more types
            key="build_type_local"
        )

        st.subheader(f"Download {build_model_type} Model")
        # Model ID Input (allow searching/pasting)
        hf_model_id = st.text_input(
            "Hugging Face Model ID",
            placeholder=f"e.g., {'google/gemma-2-9b-it' if build_model_type == 'Causal LM' else 'llava-hf/llava-1.5-7b-hf' if build_model_type == 'Vision/Multimodal' else 'microsoft/trocr-base-handwritten' if build_model_type == 'OCR' else 'stabilityai/stable-diffusion-xl-base-1.0'}",
            key="build_hf_model_id"
        )
        local_model_name = st.text_input(
            "Local Name for this Model",
            value=f"{build_model_type.split('/')[0].lower()}_{os.path.basename(hf_model_id).replace('.','') if hf_model_id else 'model'}",
            key="build_local_name",
            help="A unique name to identify this model locally."
        )

        # Add a note about token requirements for gated models
        st.info("Private or gated models require a valid Hugging Face token (set via HF_TOKEN env var or the Custom Key in sidebar API settings).")

        if st.button(f"Download & Save '{hf_model_id}' Locally", key="build_download_button", disabled=not hf_model_id or not local_model_name):
            # Validate local name uniqueness
            if local_model_name in [os.path.basename(p) for p in st.session_state.get('local_models', {})]:
                 st.error(f"A local model named '{local_model_name}' already exists. Choose a different name.")
            else:
                model_type_map = {
                    "Causal LM": "causal", "Vision/Multimodal": "vision", "OCR": "ocr", "Diffusion": "diffusion"
                }
                model_type_short = model_type_map.get(build_model_type, "unknown")

                config = LocalModelConfig(
                    name=local_model_name,
                    hf_id=hf_model_id,
                    model_type=model_type_short
                )
                save_path = config.get_full_path()
                os.makedirs(os.path.dirname(save_path), exist_ok=True)

                st.info(f"Attempting to download '{hf_model_id}' to '{save_path}'...")
                progress_bar_build = st.progress(0)
                status_text_build = st.empty()
                token_build = st.session_state.hf_custom_key or HF_TOKEN or None

                try:
                    if build_model_type == "Diffusion":
                         # Use Diffusers library download
                         if not _diffusers_available: raise ImportError("Diffusers library required for diffusion models.")
                         # Diffusers downloads directly, no explicit save needed after load typically
                         status_text_build.text("Downloading diffusion model pipeline...")
                         pipeline_obj = StableDiffusionPipeline.from_pretrained(hf_model_id, token=token_build)
                         status_text_build.text("Saving diffusion model pipeline...")
                         pipeline_obj.save_pretrained(save_path)
                         # Store info, but maybe not the full pipeline object in session state due to size
                         st.session_state.local_models[save_path] = {'type': 'diffusion', 'hf_id': hf_model_id, 'model':None, 'tokenizer':None} # Mark as downloaded
                         st.success(f"Diffusion model '{hf_model_id}' downloaded and saved to {save_path}")

                    else:
                         # Use Transformers library download
                         status_text_build.text("Downloading model components...")
                         # Determine AutoModel class based on type (can be refined)
                         if model_type_short == 'causal':
                             model_class = AutoModelForCausalLM
                             tokenizer_class = AutoTokenizer
                             processor_class = None
                         elif model_type_short == 'vision':
                             model_class = AutoModelForVision2Seq # Common for VQA/Captioning
                             processor_class = AutoProcessor # Handles image+text
                             tokenizer_class = None # Usually part of processor
                         elif model_type_short == 'ocr':
                             model_class = AutoModelForVision2Seq # TrOCR uses this
                             processor_class = AutoProcessor
                             tokenizer_class = None
                         else:
                             raise ValueError(f"Unknown model type for downloading: {model_type_short}")

                         # Download and save model
                         model_obj = model_class.from_pretrained(hf_model_id, token=token_build)
                         model_obj.save_pretrained(save_path)
                         status_text_build.text(f"Model saved. Downloading processor/tokenizer...")

                         # Download and save tokenizer/processor
                         if processor_class:
                             processor_obj = processor_class.from_pretrained(hf_model_id, token=token_build)
                             processor_obj.save_pretrained(save_path)
                             tokenizer_obj = getattr(processor_obj, 'tokenizer', None) # Get tokenizer from processor if exists
                         elif tokenizer_class:
                             tokenizer_obj = tokenizer_class.from_pretrained(hf_model_id, token=token_build)
                             tokenizer_obj.save_pretrained(save_path)
                             processor_obj = None # No separate processor
                         else: # Should not happen with current logic
                              tokenizer_obj = None
                              processor_obj = None

                         # --- Load into memory and store in session state ---
                         # This might consume significant memory! Consider loading on demand instead.
                         status_text_build.text(f"Loading '{local_model_name}' into memory...")
                         device = "cuda" if torch.cuda.is_available() else "cpu"

                         reloaded_model = model_class.from_pretrained(save_path).to(device)
                         reloaded_processor = processor_class.from_pretrained(save_path) if processor_class else None
                         reloaded_tokenizer = tokenizer_class.from_pretrained(save_path) if tokenizer_class and not reloaded_processor else getattr(reloaded_processor, 'tokenizer', None)

                         st.session_state.local_models[save_path] = {
                             'type': model_type_short,
                             'hf_id': hf_model_id,
                             'model': reloaded_model,
                             'tokenizer': reloaded_tokenizer,
                             'processor': reloaded_processor, # Store processor if it exists
                         }
                         st.success(f"{build_model_type} model '{hf_model_id}' downloaded to {save_path} and loaded into memory ({device}).")
                         # Optionally select the newly loaded model
                         st.session_state.selected_local_model_path = save_path


                except (RepositoryNotFoundError, GatedRepoError) as e:
                     st.error(f"Download failed: Repository not found or requires specific access/token. Check Model ID and your HF token. Error: {e}")
                     logger.error(f"Download failed for {hf_model_id}: {e}")
                     if os.path.exists(save_path): shutil.rmtree(save_path) # Clean up partial download
                except ImportError as e:
                     st.error(f"Download failed: Required library missing. {e}")
                     logger.error(f"ImportError during download of {hf_model_id}: {e}")
                except Exception as e:
                     st.error(f"An unexpected error occurred during download: {e}")
                     logger.error(f"Download failed for {hf_model_id}: {e}")
                     if os.path.exists(save_path): shutil.rmtree(save_path) # Clean up

                finally:
                     progress_bar_build.progress(1.0)
                     status_text_build.empty()

        st.subheader("Manage Local Models")
        loaded_model_paths = list(st.session_state.get('local_models', {}).keys())
        if not loaded_model_paths:
            st.info("No models downloaded yet.")
        else:
            models_df_data = []
            for path, data in st.session_state.local_models.items():
                 models_df_data.append({
                      "Local Name": os.path.basename(path),
                      "Type": data.get('type', 'N/A'),
                      "HF ID": data.get('hf_id', 'N/A'),
                      "Loaded": "Yes" if data.get('model') else "No (Info only)",
                      "Path": path
                 })
            models_df = pd.DataFrame(models_df_data)
            st.dataframe(models_df, use_container_width=True, hide_index=True, column_order=["Local Name", "Type", "HF ID", "Loaded"])

            model_to_delete = st.selectbox("Select model to delete", [""] + [os.path.basename(p) for p in loaded_model_paths], key="delete_model_select")
            if model_to_delete and st.button(f"Delete Local Model '{model_to_delete}'", type="primary"):
                path_to_delete = next((p for p in loaded_model_paths if os.path.basename(p) == model_to_delete), None)
                if path_to_delete:
                    try:
                        # Remove from session state first
                        del st.session_state.local_models[path_to_delete]
                        if st.session_state.selected_local_model_path == path_to_delete:
                            st.session_state.selected_local_model_path = None
                        # Delete from disk
                        if os.path.exists(path_to_delete):
                            shutil.rmtree(path_to_delete)
                        st.success(f"Deleted model '{model_to_delete}' and its files.")
                        logger.info(f"Deleted local model: {path_to_delete}")
                    except Exception as e:
                        st.error(f"Failed to delete model '{model_to_delete}': {e}")
                        logger.error(f"Failed to delete model {path_to_delete}: {e}")


# --- Tab 5: PDF Process (HF) ---
with tabs[4]:
    st.header("PDF Process with HF Models 📄")
    st.markdown("Upload PDFs, view pages, and extract text using selected HF models (API or Local).")

    # Inference Source Selection
    pdf_use_api = st.radio(
        "Choose Processing Method",
        ["Hugging Face API", "Loaded Local Model"],
        key="pdf_process_source",
        horizontal=True,
        help="API uses settings from sidebar. Local uses the selected local model (if suitable)."
    )

    if pdf_use_api == "Hugging Face API":
        st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}")
    else:
        if st.session_state.selected_local_model_path:
            st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else:
            st.warning("No local model selected. Please select one in the sidebar.")

    uploaded_pdfs_process_hf = st.file_uploader("Upload PDF files to process", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader_hf")

    if uploaded_pdfs_process_hf:
        # Simplified: Process only the first page for demonstration
        process_all_pages_pdf = st.checkbox("Process All Pages (can be slow/expensive)", value=False, key="pdf_process_all_hf")
        pdf_prompt = st.text_area("Prompt for PDF Page Processing", "Extract the text content from this page.", key="pdf_process_prompt_hf")

        if st.button("Process Uploaded PDFs with HF", key="process_uploaded_pdfs_hf"):
            if pdf_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path:
                 st.error("Cannot process locally: No local model selected.")
            else:
                combined_text_output_hf = f"# HF PDF Processing Results ({'API' if pdf_use_api else 'Local'})\n\n"
                total_pages_processed_hf = 0
                output_placeholder_hf = st.container()

                for pdf_file in uploaded_pdfs_process_hf:
                    output_placeholder_hf.markdown(f"--- \n### Processing: {pdf_file.name}")
                    # Read PDF bytes
                    pdf_bytes = pdf_file.read()
                    try:
                        doc = fitz.open("pdf", pdf_bytes) # Open from bytes
                        num_pages = len(doc)
                        pages_to_process = range(num_pages) if process_all_pages_pdf else range(min(1, num_pages)) # Limit to 1 unless checked

                        output_placeholder_hf.info(f"Processing {len(pages_to_process)} of {num_pages} pages...")
                        doc_text = f"## File: {pdf_file.name}\n\n"

                        for i in pages_to_process:
                            page_placeholder = output_placeholder_hf.empty()
                            page_placeholder.info(f"Processing Page {i + 1}/{num_pages}...")
                            page = doc[i]
                            pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

                            # Display image and process
                            cols_pdf = output_placeholder_hf.columns(2)
                            cols_pdf[0].image(img, caption=f"Page {i+1}", use_container_width=True)
                            with cols_pdf[1]:
                                 # Use the new image processing function
                                 # NOTE: This relies on the process_image_hf implementation
                                 # which is currently basic/placeholder for local models.
                                 with st.spinner("Processing page with HF model..."):
                                     hf_text = process_image_hf(img, pdf_prompt, use_api=pdf_use_api)
                                 st.text_area(f"Result (Page {i+1})", hf_text, height=250, key=f"pdf_hf_out_{pdf_file.name}_{i}")

                            doc_text += f"### Page {i + 1}\n\n{hf_text}\n\n---\n\n"
                            total_pages_processed_hf += 1
                            page_placeholder.empty() # Clear status message

                        combined_text_output_hf += doc_text
                        doc.close()

                    except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err:
                        output_placeholder_hf.error(f"Error opening PDF {pdf_file.name}: {pdf_err}. Skipping.")
                    except Exception as e:
                        output_placeholder_hf.error(f"Error processing {pdf_file.name}: {str(e)}")

                if total_pages_processed_hf > 0:
                    st.markdown("--- \n### Combined Processing Results")
                    st.text_area("Full Output", combined_text_output_hf, height=400, key="combined_pdf_hf_output")
                    output_filename_pdf_hf = generate_filename("hf_processed_pdfs", "md")
                    try:
                        with open(output_filename_pdf_hf, "w", encoding="utf-8") as f: f.write(combined_text_output_hf)
                        st.success(f"Combined output saved to {output_filename_pdf_hf}")
                        st.markdown(get_download_link(output_filename_pdf_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True)
                        st.session_state['asset_checkboxes'][output_filename_pdf_hf] = False; update_gallery()
                    except IOError as e: st.error(f"Failed to save combined output file: {e}")


# --- Tab 6: Image Process (HF) ---
with tabs[5]:
    st.header("Image Process with HF Models 🖼️")
    st.markdown("Upload images and process them using selected HF models (API or Local).")

    img_use_api = st.radio(
        "Choose Processing Method",
        ["Hugging Face API", "Loaded Local Model"],
        key="img_process_source_hf",
        horizontal=True
    )
    if img_use_api == "Hugging Face API":
        st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}")
    else:
        if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else: st.warning("No local model selected.")

    img_prompt_hf = st.text_area("Prompt for Image Processing", "Describe this image in detail.", key="img_process_prompt_hf")
    uploaded_images_process_hf = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader_hf")

    if uploaded_images_process_hf:
        if st.button("Process Uploaded Images with HF", key="process_images_hf"):
             if img_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path:
                 st.error("Cannot process locally: No local model selected.")
             else:
                combined_img_text_hf = f"# HF Image Processing Results ({'API' if img_use_api else 'Local'})\n\n**Prompt:** {img_prompt_hf}\n\n---\n\n"
                images_processed_hf = 0
                output_img_placeholder_hf = st.container()

                for img_file in uploaded_images_process_hf:
                    output_img_placeholder_hf.markdown(f"### Processing: {img_file.name}")
                    try:
                        img = Image.open(img_file)
                        cols_img_hf = output_img_placeholder_hf.columns(2)
                        cols_img_hf[0].image(img, caption=f"Input: {img_file.name}", use_container_width=True)
                        with cols_img_hf[1], st.spinner("Processing image with HF model..."):
                             # Use the new image processing function
                             hf_img_text = process_image_hf(img, img_prompt_hf, use_api=img_use_api)
                             st.text_area(f"Result", hf_img_text, height=300, key=f"img_hf_out_{img_file.name}")

                        combined_img_text_hf += f"## Image: {img_file.name}\n\n{hf_img_text}\n\n---\n\n"
                        images_processed_hf += 1

                    except UnidentifiedImageError: output_img_placeholder_hf.error(f"Invalid Image: {img_file.name}. Skipping.")
                    except Exception as e: output_img_placeholder_hf.error(f"Error processing {img_file.name}: {str(e)}")

                if images_processed_hf > 0:
                    st.markdown("--- \n### Combined Processing Results")
                    st.text_area("Full Output", combined_img_text_hf, height=400, key="combined_img_hf_output")
                    output_filename_img_hf = generate_filename("hf_processed_images", "md")
                    try:
                        with open(output_filename_img_hf, "w", encoding="utf-8") as f: f.write(combined_img_text_hf)
                        st.success(f"Combined output saved to {output_filename_img_hf}")
                        st.markdown(get_download_link(output_filename_img_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True)
                        st.session_state['asset_checkboxes'][output_filename_img_hf] = False; update_gallery()
                    except IOError as e: st.error(f"Failed to save combined output file: {e}")


# --- Tab 7: Text Process (HF) ---
with tabs[6]:
    st.header("Text Process with HF Models 📝")
    st.markdown("Process Markdown (.md) or Text (.txt) files using selected HF models (API or Local).")

    text_use_api = st.radio(
        "Choose Processing Method",
        ["Hugging Face API", "Loaded Local Model"],
        key="text_process_source_hf",
        horizontal=True
    )
    if text_use_api == "Hugging Face API":
        st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}")
    else:
        if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
        else: st.warning("No local model selected.")

    text_files_hf = get_gallery_files(['md', 'txt'])
    if not text_files_hf:
         st.warning("No .md or .txt files in gallery to process.")
    else:
        selected_text_file_hf = st.selectbox(
             "Select Text/MD File to Process",
             options=[""] + text_files_hf,
             format_func=lambda x: os.path.basename(x) if x else "Select a file...",
             key="text_process_select_hf"
        )

        if selected_text_file_hf:
             st.write(f"Selected: {os.path.basename(selected_text_file_hf)}")
             try:
                 with open(selected_text_file_hf, "r", encoding="utf-8", errors='ignore') as f:
                     content_text_hf = f.read()
                 st.text_area("File Content Preview", content_text_hf[:1000] + ("..." if len(content_text_hf) > 1000 else ""), height=200, key="text_content_preview_hf")

                 prompt_text_hf = st.text_area(
                     "Enter Prompt for this File",
                     "Summarize the key points of this text.",
                     key="text_individual_prompt_hf"
                 )

                 if st.button(f"Process '{os.path.basename(selected_text_file_hf)}' with HF", key=f"process_text_hf_btn"):
                     if text_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path:
                         st.error("Cannot process locally: No local model selected.")
                     else:
                        with st.spinner("Processing text with HF model..."):
                            result_text_processed = process_text_hf(content_text_hf, prompt_text_hf, use_api=text_use_api)

                        st.markdown("### Processing Result")
                        st.markdown(result_text_processed) # Display result

                        output_filename_text_hf = generate_filename(f"hf_processed_{os.path.splitext(os.path.basename(selected_text_file_hf))[0]}", "md")
                        try:
                            with open(output_filename_text_hf, "w", encoding="utf-8") as f: f.write(result_text_processed)
                            st.success(f"Result saved to {output_filename_text_hf}")
                            st.markdown(get_download_link(output_filename_text_hf, "text/markdown", "Download Result MD"), unsafe_allow_html=True)
                            st.session_state['asset_checkboxes'][output_filename_text_hf] = False; update_gallery()
                        except IOError as e: st.error(f"Failed to save result file: {e}")

             except FileNotFoundError: st.error("Selected file not found.")
             except Exception as e: st.error(f"Error reading file: {e}")


# --- Tab 8: Test OCR (HF) ---
with tabs[7]:
    st.header("Test OCR with HF Models 🔍")
    st.markdown("Select an image/PDF and run OCR using HF models (API or Local - requires suitable local model).")

    ocr_use_api = st.radio(
        "Choose OCR Method",
        ["Hugging Face API (Basic Captioning/OCR)", "Loaded Local OCR Model"],
        key="ocr_source_hf",
        horizontal=True,
        help="API uses basic image-to-text. Local requires a dedicated OCR model (e.g., TrOCR) to be loaded."
    )
    if ocr_use_api == "Loaded Local OCR Model":
         if st.session_state.selected_local_model_path:
              model_type = st.session_state.local_models.get(st.session_state.selected_local_model_path,{}).get('type')
              if model_type != 'ocr':
                   st.warning(f"Selected local model ({os.path.basename(st.session_state.selected_local_model_path)}) is type '{model_type}', not 'ocr'. Results may be poor.")
              else:
                   st.info(f"Using Local OCR Model: {os.path.basename(st.session_state.selected_local_model_path)}")
         else: st.warning("No local model selected.")

    gallery_files_ocr_hf = get_gallery_files(['png', 'jpg', 'jpeg', 'pdf'])
    if not gallery_files_ocr_hf:
        st.warning("No images or PDFs in gallery.")
    else:
        selected_file_ocr_hf = st.selectbox(
            "Select Image or PDF from Gallery for OCR",
            options=[""] + gallery_files_ocr_hf,
            format_func=lambda x: os.path.basename(x) if x else "Select a file...",
            key="ocr_select_file_hf"
        )

        if selected_file_ocr_hf:
            st.write(f"Selected: {os.path.basename(selected_file_ocr_hf)}")
            file_ext_ocr_hf = os.path.splitext(selected_file_ocr_hf)[1].lower()
            image_to_ocr_hf = None; page_info_hf = ""

            try:
                if file_ext_ocr_hf in ['.png', '.jpg', '.jpeg']: image_to_ocr_hf = Image.open(selected_file_ocr_hf)
                elif file_ext_ocr_hf == '.pdf':
                    doc = fitz.open(selected_file_ocr_hf)
                    if len(doc) > 0: pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image_to_ocr_hf = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); page_info_hf = " (Page 1)"
                    else: st.warning("Selected PDF is empty.")
                    doc.close()

                if image_to_ocr_hf:
                    st.image(image_to_ocr_hf, caption=f"Image for OCR{page_info_hf}", use_container_width=True)
                    if st.button("Run HF OCR on this Image 🚀", key="ocr_run_button_hf"):
                        if ocr_use_api == "Loaded Local OCR Model" and not st.session_state.selected_local_model_path:
                             st.error("Cannot run locally: No local model selected.")
                        else:
                            output_ocr_file_hf = generate_filename(f"hf_ocr_{os.path.splitext(os.path.basename(selected_file_ocr_hf))[0]}", "txt")
                            st.session_state['processing']['ocr'] = True
                            with st.spinner("Performing OCR with HF model..."):
                                 ocr_result_hf = asyncio.run(process_hf_ocr(image_to_ocr_hf, output_ocr_file_hf, use_api=ocr_use_api))
                            st.session_state['processing']['ocr'] = False

                            st.text_area("OCR Result", ocr_result_hf, height=300, key="ocr_result_display_hf")
                            if ocr_result_hf and not ocr_result_hf.startswith("Error") and not ocr_result_hf.startswith("["):
                                entry = f"HF OCR: {selected_file_ocr_hf}{page_info_hf} -> {output_ocr_file_hf}"
                                st.session_state['history'].append(entry)
                                if len(ocr_result_hf) > 5: # Minimal check
                                    st.success(f"OCR output saved to {output_ocr_file_hf}")
                                    st.markdown(get_download_link(output_ocr_file_hf, "text/plain", "Download OCR Text"), unsafe_allow_html=True)
                                    st.session_state['asset_checkboxes'][output_ocr_file_hf] = False; update_gallery()
                                else: st.warning("OCR output seems short/empty.")
                            else: st.error(f"OCR failed. {ocr_result_hf}")

            except Exception as e: st.error(f"Error loading file for OCR: {e}")

# --- Tab 9: Test Image Gen (Diffusers) ---
with tabs[8]:
    st.header("Test Image Generation (Diffusers) 🎨")
    st.markdown("Generate images using Stable Diffusion models loaded locally via the Diffusers library.")

    if not _diffusers_available:
         st.error("Diffusers library is required for image generation.")
    else:
        # Select from locally downloaded *diffusion* models
        local_diffusion_paths = get_local_model_paths("diffusion")
        if not local_diffusion_paths:
             st.warning("No local diffusion models found. Download one using the 'Build Titan' tab.")
             selected_diffusion_model_path = None
        else:
            selected_diffusion_model_path = st.selectbox(
                 "Select Local Diffusion Model",
                 options=[""] + local_diffusion_paths,
                 format_func=lambda x: os.path.basename(x) if x else "Select...",
                 key="imggen_diffusion_model_select"
            )

        prompt_imggen_diff = st.text_area("Image Generation Prompt", "A photorealistic cat wearing sunglasses, studio lighting", key="imggen_prompt_diff")
        neg_prompt_imggen_diff = st.text_area("Negative Prompt (Optional)", "ugly, deformed, blurry, low quality", key="imggen_neg_prompt_diff")
        steps_imggen_diff = st.slider("Inference Steps", 10, 100, 25, key="imggen_steps")
        guidance_imggen_diff = st.slider("Guidance Scale", 1.0, 20.0, 7.5, step=0.5, key="imggen_guidance")

        if st.button("Generate Image 🚀", key="imggen_run_button_diff", disabled=not selected_diffusion_model_path):
            if not prompt_imggen_diff: st.warning("Please enter a prompt.")
            else:
                 status_imggen = st.empty()
                 try:
                     # Load pipeline from saved path on demand
                     status_imggen.info(f"Loading diffusion pipeline: {os.path.basename(selected_diffusion_model_path)}...")
                     # Determine device
                     device = "cuda" if torch.cuda.is_available() else "cpu"
                     dtype = torch.float16 if torch.cuda.is_available() else torch.float32 # Use float16 on GPU if available
                     pipe = StableDiffusionPipeline.from_pretrained(selected_diffusion_model_path, torch_dtype=dtype).to(device)
                     pipe.safety_checker = None # Optional: Disable safety checker if needed

                     status_imggen.info(f"Generating image on {device} ({dtype})...")
                     start_gen_time = time.time()

                     # Generate using the pipeline
                     gen_output = pipe(
                          prompt=prompt_imggen_diff,
                          negative_prompt=neg_prompt_imggen_diff if neg_prompt_imggen_diff else None,
                          num_inference_steps=steps_imggen_diff,
                          guidance_scale=guidance_imggen_diff,
                          # Add seed if desired: generator=torch.Generator(device=device).manual_seed(your_seed)
                     )
                     gen_image = gen_output.images[0]

                     elapsed_gen = int(time.time() - start_gen_time)
                     status_imggen.success(f"Image generated in {elapsed_gen}s!")

                     # Save and display
                     output_imggen_file_diff = generate_filename("diffusion_gen", "png")
                     gen_image.save(output_imggen_file_diff)
                     st.image(gen_image, caption=f"Generated: {output_imggen_file_diff}", use_container_width=True)
                     st.markdown(get_download_link(output_imggen_file_diff, "image/png", "Download Generated Image"), unsafe_allow_html=True)
                     st.session_state['asset_checkboxes'][output_imggen_file_diff] = False; update_gallery()
                     st.session_state['history'].append(f"Diffusion Gen: '{prompt_imggen_diff[:30]}...' -> {output_imggen_file_diff}")

                 except ImportError: st.error("Diffusers or Torch library not found.")
                 except Exception as e:
                     st.error(f"Image generation failed: {e}")
                     logger.error(f"Diffusion generation failed for {selected_diffusion_model_path}: {e}")
                 finally:
                     # Clear pipeline from memory? (Optional, depends on memory usage)
                     if 'pipe' in locals(): del pipe; torch.cuda.empty_cache() if torch.cuda.is_available() else None


# --- Tab 10: Character Editor (Keep from previous merge) ---
with tabs[9]:
    # ... (Code from previous merge for this tab) ...
    st.header("Character Editor 🧑‍🎨")
    st.subheader("Create Your Character")
    load_characters(); existing_char_names = [c['name'] for c in st.session_state.get('characters', [])]
    form_key = f"character_form_{st.session_state.get('char_form_reset_key', 0)}"
    with st.form(key=form_key):
        st.markdown("**Create New Character**")
        if st.form_submit_button("Randomize Content 🎲"): st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1; st.rerun()
        rand_name, rand_gender, rand_intro, rand_greeting = randomize_character_content()
        name_char = st.text_input("Name (3-25 chars...)", value=rand_name, max_chars=25, key="char_name_input")
        gender_char = st.radio("Gender", ["Male", "Female"], index=["Male", "Female"].index(rand_gender), key="char_gender_radio")
        intro_char = st.text_area("Intro (Public description)", value=rand_intro, max_chars=300, height=100, key="char_intro_area")
        greeting_char = st.text_area("Greeting (First message)", value=rand_greeting, max_chars=300, height=100, key="char_greeting_area")
        tags_char = st.text_input("Tags (comma-separated)", "OC, friendly", key="char_tags_input")
        submitted = st.form_submit_button("Create Character ✨")
        if submitted:
            error = False
            if not (3 <= len(name_char) <= 25): st.error("Name must be 3-25 characters."); error = True
            if not re.match(r'^[a-zA-Z0-9 _-]+$', name_char): st.error("Name contains invalid characters."); error = True
            if name_char in existing_char_names: st.error(f"Name '{name_char}' already exists!"); error = True
            if not intro_char or not greeting_char: st.error("Intro/Greeting cannot be empty."); error = True
            if not error:
                tag_list = [tag.strip() for tag in tags_char.split(',') if tag.strip()]
                character_data = {"name": name_char, "gender": gender_char, "intro": intro_char, "greeting": greeting_char, "created_at": datetime.now(pytz.timezone("US/Central")).strftime('%Y-%m-%d %H:%M:%S %Z'), "tags": tag_list}
                if save_character(character_data):
                    st.success(f"Character '{name_char}' created!"); st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1; st.rerun()

# --- Tab 11: Character Gallery (Keep from previous merge) ---
with tabs[10]:
    # ... (Code from previous merge for this tab) ...
    st.header("Character Gallery 🖼️")
    load_characters(); characters_list = st.session_state.get('characters', [])
    if not characters_list: st.warning("No characters created yet.")
    else:
        st.subheader(f"Your Characters ({len(characters_list)})")
        search_term = st.text_input("Search Characters by Name", key="char_gallery_search")
        if search_term: characters_list = [c for c in characters_list if search_term.lower() in c['name'].lower()]
        cols_char_gallery = st.columns(3); chars_to_delete = []
        for idx, char in enumerate(characters_list):
            with cols_char_gallery[idx % 3], st.container(border=True):
                st.markdown(f"**{char['name']}**"); st.caption(f"Gender: {char.get('gender', 'N/A')}")
                st.markdown("**Intro:**"); st.markdown(f"> {char.get('intro', '')}")
                st.markdown("**Greeting:**"); st.markdown(f"> {char.get('greeting', '')}")
                st.caption(f"Tags: {', '.join(char.get('tags', ['N/A']))}"); st.caption(f"Created: {char.get('created_at', 'N/A')}")
                delete_key_char = f"delete_char_{char['name']}_{idx}";
                if st.button(f"Delete {char['name']}", key=delete_key_char, type="primary"): chars_to_delete.append(char['name'])
        if chars_to_delete:
             current_characters = st.session_state.get('characters', []); updated_characters = [c for c in current_characters if c['name'] not in chars_to_delete]
             st.session_state['characters'] = updated_characters
             try:
                 with open("characters.json", "w", encoding='utf-8') as f: json.dump(updated_characters, f, indent=2)
                 logger.info(f"Deleted characters: {', '.join(chars_to_delete)}"); st.success(f"Deleted characters: {', '.join(chars_to_delete)}"); st.rerun()
             except IOError as e: logger.error(f"Failed to save characters.json after deletion: {e}"); st.error("Failed to update character file.")

# --- Footer and Persistent Sidebar Elements ------------

# Update Sidebar Gallery (Call this at the end to reflect all changes)
update_gallery()

# Action Logs in Sidebar
st.sidebar.subheader("Action Logs 📜")
log_expander = st.sidebar.expander("View Logs", expanded=False)
with log_expander:
    log_text = "\n".join([f"{record.asctime} - {record.levelname} - {record.message}" for record in log_records[-20:]])
    st.code(log_text, language='log')

# History in Sidebar
st.sidebar.subheader("Session History 📜")
history_expander = st.sidebar.expander("View History", expanded=False)
with history_expander:
     for entry in reversed(st.session_state.get("history", [])):
         if entry: history_expander.write(f"- {entry}")

st.sidebar.markdown("---")
st.sidebar.info("Using Hugging Face models for AI tasks.")
st.sidebar.caption("App Modified by AI Assistant")