Update app.py

app.py

@@ -1,115 +1,939 @@
-# app.py (streamlined and modularized for clarity and maintainability)
-import streamlit as st
+import io
 import os
-import glob
+import re
 import base64
-import pandas as pd
-import fitz
-from PIL import Image
-from io import BytesIO
+import glob
+import logging
+import random
+import shutil
+import time
+import zipfile
+import json
+import asyncio
+import aiofiles
+import toml
 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
 from reportlab.pdfgen import canvas
 from reportlab.lib.utils import ImageReader
-from markdown2 import markdown
+from reportlab.lib.pagesizes import letter
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.lib.enums import TA_JUSTIFY
+import fitz
+import requests
+try:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, pipeline
+    _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
+    if _transformers_available:
+        st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.")
+try:
+    from openai import OpenAI
+    _openai_available = True
+except ImportError:
+    _openai_available = False
+    st.sidebar.warning("OpenAI library not found. OpenAI model features disabled.")
+from huggingface_hub import InferenceClient, HfApi, list_models
+from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError
 
-# --- Config ---
+# --- App Configuration ---
 st.set_page_config(
-    page_title="Vision & Layout Titans 🚀",
+    page_title="Vision & Layout Titans 🚀🖼️",
     page_icon="🤖",
-    layout="wide"
+    layout="wide",
+    initial_sidebar_state="expanded",
+    menu_items={
+        'Get Help': 'https://huggingface.co/docs',
+        'Report a Bug': None,
+        'About': "Combined App: Image/MD->PDF Layout + AI-Powered Tools 🌌"
+    }
 )
 
+# --- Secrets Management ---
+try:
+    secrets = toml.load(".streamlit/secrets.toml") if os.path.exists(".streamlit/secrets.toml") else {}
+    HF_TOKEN = secrets.get("HF_TOKEN", os.getenv("HF_TOKEN", ""))
+    OPENAI_API_KEY = secrets.get("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY", ""))
+except Exception as e:
+    st.error(f"Error loading secrets: {e}")
+    HF_TOKEN = os.getenv("HF_TOKEN", "")
+    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
+
+if not HF_TOKEN:
+    st.sidebar.warning("Hugging Face token not found in secrets or environment. Some features may be limited.")
+if not OPENAI_API_KEY and _openai_available:
+    st.sidebar.warning("OpenAI API key not found in secrets or environment. OpenAI features disabled.")
+
+# --- 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())
+
+# --- Model Initialization ---
+DEFAULT_PROVIDER = "hf-inference"
+FEATURED_MODELS_LIST = [
+    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    "google/gemma-2-9b-it",
+    "Qwen/Qwen2-7B-Instruct",
+    "microsoft/Phi-3-mini-4k-instruct",
+    "HuggingFaceH4/zephyr-7b-beta",
+    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+    "HuggingFaceTB/SmolLM-1.7B-Instruct"
+]
+VISION_MODELS_LIST = [
+    "Salesforce/blip-image-captioning-large",
+    "microsoft/trocr-large-handwritten",
+    "llava-hf/llava-1.5-7b-hf",
+    "google/vit-base-patch16-224"
+]
+DIFFUSION_MODELS_LIST = [
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "runwayml/stable-diffusion-v1-5",
+    "OFA-Sys/small-stable-diffusion-v0"
+]
+OPENAI_MODELS_LIST = [
+    "gpt-4o",
+    "gpt-4-turbo",
+    "gpt-3.5-turbo",
+    "text-davinci-003"
+]
+st.session_state.setdefault('local_models', {})
+st.session_state.setdefault('hf_inference_client', None)
+st.session_state.setdefault('openai_client', None)
+if _openai_available and OPENAI_API_KEY:
+    try:
+        st.session_state['openai_client'] = OpenAI(api_key=OPENAI_API_KEY)
+        logger.info("OpenAI client initialized successfully.")
+    except Exception as e:
+        st.error(f"Failed to initialize OpenAI client: {e}")
+        logger.error(f"OpenAI client initialization failed: {e}")
+        st.session_state['openai_client'] = None
+
+# --- Session State Initialization ---
+st.session_state.setdefault('layout_snapshots', [])
+st.session_state.setdefault('layout_new_uploads', [])
+st.session_state.setdefault('history', [])
+st.session_state.setdefault('processing', {})
+st.session_state.setdefault('asset_checkboxes', {'image': {}, 'md': {}, 'pdf': {}})
+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)
+st.session_state.setdefault('gallery_size', 10)
+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])
+st.session_state.setdefault('hf_custom_api_model', "")
+st.session_state.setdefault('openai_selected_model', OPENAI_MODELS_LIST[0] if _openai_available else "")
+st.session_state.setdefault('selected_local_model_path', None)
+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)
+if 'asset_gallery_container' not in st.session_state:
+    st.session_state['asset_gallery_container'] = {'image': st.sidebar.empty(), 'md': st.sidebar.empty(), 'pdf': st.sidebar.empty()}
+
+# --- Dataclasses ---
+@dataclass
+class LocalModelConfig:
+    name: str
+    hf_id: str
+    model_type: str
+    size_category: str = "unknown"
+    domain: Optional[str] = None
+    local_path: str = field(init=False)
+    def __post_init__(self):
+        type_folder = f"{self.model_type}_models"
+        safe_name = re.sub(r'[^\w\-]+', '_', self.name)
+        self.local_path = os.path.join(type_folder, safe_name)
+    def get_full_path(self):
+        return os.path.abspath(self.local_path)
+
+@dataclass
+class DiffusionConfig:
+    name: str
+    base_model: str
+    size: str
+    domain: Optional[str] = None
+    @property
+    def model_path(self):
+        return f"diffusion_models/{self.name}"
+
 # --- Helper Functions ---
-def get_files(exts):
-    files = []
-    for ext in exts:
-        files.extend(glob.glob(f'*.{ext}'))
-    return sorted([f for f in files if f.lower() != "readme.md"])
-
-def image_to_pdf(images, md_files):
-    buffer = BytesIO()
-    c = canvas.Canvas(buffer)
-    
-    for img_path in images:
-        img = Image.open(img_path)
-        width, height = img.size
-        c.setPageSize((width, height))
-        c.drawImage(ImageReader(img), 0, 0, width, height)
-        c.showPage()
-
-    for md_path in md_files:
-        with open(md_path, 'r', encoding='utf-8') as f:
-            md_content = f.read()
-            html = markdown(md_content)
-        c.setPageSize((595, 842))  # A4 size
-        c.setFont("Helvetica", 10)
-        c.drawString(50, 800, md_content[:1000])  # Simplified render
-        c.showPage()
-
-    c.save()
-    buffer.seek(0)
-    return buffer
-
-def render_pdf_gallery(pdf_files):
-    for pdf_file in pdf_files:
-        doc = fitz.open(pdf_file)
-        page = doc.load_page(0)
-        pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        st.image(img, caption=os.path.basename(pdf_file))
-
-# --- Sidebar Management ---
-st.sidebar.header("Content Management")
-
-image_files = get_files(['png', 'jpg', 'jpeg'])
-pdf_files = get_files(['pdf'])
-md_files = get_files(['md'])
-
-selected_images = st.sidebar.multiselect("Select Images", image_files)
-selected_md_files = st.sidebar.multiselect("Select Markdown Files", md_files)
-
-# PDF Gallery
-if st.sidebar.checkbox("Show PDF Gallery"):
-    render_pdf_gallery(pdf_files)
-
-# PDF Generation with rearrangement
-st.sidebar.subheader("Reorder Content for PDF")
-content_df = pd.DataFrame({
-    "Content": selected_images + selected_md_files,
-    "Type": ["Image"] * len(selected_images) + ["Markdown"] * len(selected_md_files),
-    "Order": range(len(selected_images) + len(selected_md_files))
-})
-
-edited_df = st.sidebar.data_editor(content_df, use_container_width=True)
-sorted_contents = edited_df.sort_values('Order')['Content'].tolist()
-
-if st.sidebar.button("Generate PDF"):
-    sorted_images = [item for item in sorted_contents if item in selected_images]
-    sorted_md_files = [item for item in sorted_contents if item in selected_md_files]
-    pdf_buffer = image_to_pdf(sorted_images, sorted_md_files)
-    st.sidebar.download_button("Download PDF", pdf_buffer, "output.pdf")
-
-# Deletion
-st.sidebar.subheader("Delete Files")
-file_to_delete = st.sidebar.selectbox("Select file to delete", image_files + pdf_files + md_files)
-if st.sidebar.button("Delete Selected File"):
-    os.remove(file_to_delete)
-    st.sidebar.success(f"Deleted {file_to_delete}")
-    st.rerun()
-
-# --- Main Page ---
-st.title("Vision & Layout Titans 🚀")
-st.markdown("### Manage, View, and Export Your Files Easily!")
-
-# Display selected images
-st.subheader("Selected Images")
-for img_path in selected_images:
-    img = Image.open(img_path)
-    st.image(img, caption=os.path.basename(img_path))
-
-# Display selected markdown
-st.subheader("Selected Markdown")
-for md_path in selected_md_files:
-    with open(md_path, 'r', encoding='utf-8') as f:
-        md_content = f.read()
-    st.markdown(md_content[:500] + '...')
+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"
+
+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"):
+    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([f for f in all_files if os.path.basename(f).lower() != 'readme.md'])
+
+def get_typed_gallery_files(file_type):
+    if file_type == 'image':
+        return get_gallery_files(('png', 'jpg', 'jpeg'))
+    elif file_type == 'md':
+        return get_gallery_files(('md',))
+    elif file_type == 'pdf':
+        return get_gallery_files(('pdf',))
+    return []
+
+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
+
+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 = min(1, len(doc)) if mode == "single" else min(2, len(doc)) if mode == "twopage" else 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 []
+
+def get_hf_client() -> Optional[InferenceClient]:
+    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.")
+        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
+    current_client = st.session_state.get('hf_inference_client')
+    needs_reinit = True
+    if current_client:
+        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}.")
+            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: {e}")
+            logger.error(f"InferenceClient initialization failed: {e}")
+            st.session_state.hf_inference_client = None
+    return st.session_state.hf_inference_client
+
+def process_text_hf(text: str, prompt: str, use_api: bool, model_id: str = None) -> str:
+    status_placeholder = st.empty()
+    start_time = time.time()
+    result_text = ""
+    params = {
+        "max_new_tokens": st.session_state.gen_max_tokens,
+        "temperature": st.session_state.gen_temperature,
+        "top_p": st.session_state.gen_top_p,
+        "repetition_penalty": st.session_state.gen_frequency_penalty + 1.0,
+    }
+    seed = st.session_state.gen_seed
+    if seed != -1:
+        params["seed"] = seed
+    system_prompt = "You are a helpful assistant. Process the following text based on the user's request."
+    full_prompt = f"{prompt}\n\n---\n\n{text}"
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": full_prompt}
+    ]
+    if use_api:
+        status_placeholder.info("Processing text using Hugging Face API...")
+        client = get_hf_client()
+        if not client:
+            return "Error: Hugging Face client not available."
+        model_id = model_id or st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
+        status_placeholder.info(f"Using API Model: {model_id}")
+        try:
+            response = client.chat_completion(
+                model=model_id,
+                messages=messages,
+                max_tokens=params['max_new_tokens'],
+                temperature=params['temperature'],
+                top_p=params['top_p'],
+            )
+            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:
+        status_placeholder.info("Processing text using local model...")
+        if not _transformers_available:
+            return "Error: Transformers library not available."
+        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."
+        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:
+            try:
+                prompt_for_model = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            except Exception:
+                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)
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            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),
+                "do_sample": True if params['temperature'] > 0.1 else False,
+                "pad_token_id": tokenizer.eos_token_id
+            }
+            with torch.no_grad():
+                outputs = model.generate(**inputs, **generate_params)
+            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
+
+def process_text_openai(text: str, prompt: str, model_id: str) -> str:
+    if not _openai_available or not st.session_state.get('openai_client'):
+        return "Error: OpenAI client not available or API key missing."
+    status_placeholder = st.empty()
+    start_time = time.time()
+    client = st.session_state['openai_client']
+    system_prompt = "You are a helpful assistant. Process the following text based on the user's request."
+    full_prompt = f"{prompt}\n\n---\n\n{text}"
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": full_prompt}
+    ]
+    status_placeholder.info(f"Processing text using OpenAI model: {model_id}...")
+    try:
+        response = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            max_tokens=st.session_state.gen_max_tokens,
+            temperature=st.session_state.gen_temperature,
+            top_p=st.session_state.gen_top_p,
+        )
+        result_text = response.choices[0].message.content or ""
+        logger.info(f"OpenAI text processing successful for model {model_id}.")
+    except Exception as e:
+        logger.error(f"OpenAI text processing failed for model {model_id}: {e}")
+        result_text = f"Error during OpenAI inference: {str(e)}"
+    elapsed = int(time.time() - start_time)
+    status_placeholder.success(f"Text processing completed in {elapsed}s.")
+    return result_text
+
+def process_image_hf(image: Image.Image, prompt: str, use_api: bool, model_id: str = None) -> str:
+    status_placeholder = st.empty()
+    start_time = time.time()
+    result_text = ""
+    if use_api:
+        status_placeholder.info("Processing image using Hugging Face API...")
+        client = get_hf_client()
+        if not client:
+            return "Error: HF client not configured."
+        buffered = BytesIO()
+        image.save(buffered, format="PNG" if image.format != 'JPEG' else 'JPEG')
+        img_bytes = buffered.getvalue()
+        model_id = model_id or "Salesforce/blip-image-captioning-large"
+        status_placeholder.info(f"Using API Image-to-Text Model: {model_id}")
+        try:
+            response_list = client.image_to_text(data=img_bytes, model=model_id)
+            if response_list and isinstance(response_list, list) and 'generated_text' in response_list[0]:
+                result_text = response_list[0]['generated_text']
+                logger.info(f"HF API image captioning successful for model {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:
+        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."
+        local_model_data = st.session_state['local_models'][model_path]
+        model_type = local_model_data.get('type')
+        if model_type == 'vision':
+            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=st.session_state.gen_max_tokens)
+                    result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+                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."
+        elif model_type == 'ocr':
+            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=st.session_state.gen_max_tokens)
+                    result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+                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."
+        else:
+            result_text = f"Error: Loaded model '{os.path.basename(model_path)}' is not a recognized vision/OCR type."
+    elapsed = int(time.time() - start_time)
+    status_placeholder.success(f"Image processing completed in {elapsed}s.")
+    return result_text
+
+def process_image_openai(image: Image.Image, prompt: str, model_id: str = "gpt-4o") -> str:
+    if not _openai_available or not st.session_state.get('openai_client'):
+        return "Error: OpenAI client not available or API key missing."
+    status_placeholder = st.empty()
+    start_time = time.time()
+    client = st.session_state['openai_client']
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_b64 = base64.b64encode(buffered.getvalue()).decode()
+    status_placeholder.info(f"Processing image using OpenAI model: {model_id}...")
+    try:
+        response = client.chat.completions.create(
+            model=model_id,
+            messages=[
+                {"role": "user", "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}}
+                ]}
+            ],
+            max_tokens=st.session_state.gen_max_tokens,
+            temperature=st.session_state.gen_temperature,
+        )
+        result_text = response.choices[0].message.content or ""
+        logger.info(f"OpenAI image processing successful for model {model_id}.")
+    except Exception as e:
+        logger.error(f"OpenAI image processing failed for model {model_id}: {e}")
+        result_text = f"Error during OpenAI image inference: {str(e)}"
+    elapsed = int(time.time() - start_time)
+    status_placeholder.success(f"Image processing completed in {elapsed}s.")
+    return result_text
+
+async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool, model_id: str = None) -> str:
+    ocr_prompt = "Extract text content from this image."
+    result = process_image_hf(image, ocr_prompt, use_api, model_id=model_id or "microsoft/trocr-large-handwritten")
+    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}]"
+    elif os.path.exists(output_file):
+        try:
+            os.remove(output_file)
+        except OSError:
+            pass
+    return result
+
+async def process_openai_ocr(image: Image.Image, output_file: str, model_id: str = "gpt-4o") -> str:
+    ocr_prompt = "Extract text content from this image."
+    result = process_image_openai(image, ocr_prompt, model_id)
+    if result and not result.startswith("Error"):
+        try:
+            async with aiofiles.open(output_file, "w", encoding='utf-8') as f:
+                await f.write(result)
+            logger.info(f"OpenAI OCR result saved to {output_file}")
+        except IOError as e:
+            logger.error(f"Failed to save OpenAI OCR output to {output_file}: {e}")
+            result += f"\n[Error saving file: {e}]"
+    elif os.path.exists(output_file):
+        try:
+            os.remove(output_file)
+        except OSError:
+            pass
+    return result
+
+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}")
+
+def clean_stem(fn: str) -> str:
+    name = os.path.splitext(os.path.basename(fn))[0]
+    name = name.replace('-', ' ').replace('_', ' ')
+    return name.strip().title()
+
+def make_image_sized_pdf(sources, is_markdown_flags):
+    if not sources:
+        st.warning("No sources provided for PDF generation.")
+        return None
+    buf = BytesIO()
+    styles = getSampleStyleSheet()
+    md_style = ParagraphStyle(
+        name='Markdown',
+        fontSize=10,
+        leading=12,
+        spaceAfter=6,
+        alignment=TA_JUSTIFY,
+        fontName='Helvetica'
+    )
+    doc = SimpleDocTemplate(buf, pagesize=letter, rightMargin=36, leftMargin=36, topMargin=36, bottomMargin=36)
+    story = []
+    try:
+        for idx, (src, is_md) in enumerate(zip(sources, is_markdown_flags), start=1):
+            status_placeholder = st.empty()
+            filename = 'page_' + str(idx)
+            status_placeholder.info(f"Adding page {idx}/{len(sources)}: {os.path.basename(str(src))}...")
+            try:
+                if is_md:
+                    with open(src, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    content = re.sub(r'!\[.*?\]\(.*?\)', '', content)
+                    paragraphs = content.split('\n\n')
+                    for para in paragraphs:
+                        if para.strip():
+                            story.append(Paragraph(para.strip(), md_style))
+                    story.append(PageBreak())
+                    status_placeholder.success(f"Added markdown page {idx}/{len(sources)}: {filename}")
+                else:
+                    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
+                        c = canvas.Canvas(BytesIO(), pagesize=(iw, ih + cap_h))
+                        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(iw / 2, cap_h / 2 + 3, caption)
+                        c.setFont('Helvetica', 8)
+                        c.setFillColorRGB(0.5, 0.5, 0.5)
+                        c.drawRightString(iw - 10, 8, f"Page {idx}")
+                        c.save()
+                        story.append(PageBreak())
+                        status_placeholder.success(f"Added image page {idx}/{len(sources)}: {filename}")
+            except Exception as e:
+                logger.error(f"Error processing source {src}: {e}")
+                status_placeholder.error(f"Error adding page {idx}: {e}")
+        doc.build(story)
+        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
+
+def update_gallery(gallery_type='image'):
+    container = st.session_state['asset_gallery_container'][gallery_type]
+    with container:
+        st.markdown(f"### {gallery_type.capitalize()} Gallery 📸")
+        files = get_typed_gallery_files(gallery_type)
+        if not files:
+            st.info(f"No {gallery_type} assets found yet.")
+            return
+        st.caption(f"Found {len(files)} assets:")
+        for idx, file in enumerate(files[:st.session_state.gallery_size]):
+            st.session_state['unique_counter'] += 1
+            unique_id = st.session_state['unique_counter']
+            item_key_base = f"{gallery_type}_gallery_item_{os.path.basename(file)}_{unique_id}"
+            basename = os.path.basename(file)
+            st.markdown(f"**{basename}**")
+            try:
+                file_ext = os.path.splitext(file)[1].lower()
+                if gallery_type == 'image' and file_ext in ['.png', '.jpg', '.jpeg']:
+                    with st.expander("Preview", expanded=False):
+                        st.image(Image.open(file), use_container_width=True)
+                elif gallery_type == 'pdf' and file_ext == '.pdf':
+                    with st.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))
+                            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 gallery_type == 'md' and file_ext == '.md':
+                    with st.expander("Preview (Start)", expanded=False):
+                        with open(file, 'r', encoding='utf-8', errors='ignore') as f:
+                            content_preview = f.read(200)
+                        st.code(content_preview + "...", language='markdown')
+                action_cols = st.columns(3)
+                with action_cols[0]:
+                    checkbox_key = f"cb_{item_key_base}"
+                    st.session_state['asset_checkboxes'][gallery_type][file] = st.checkbox(
+                        "Select",
+                        value=st.session_state['asset_checkboxes'][gallery_type].get(file, False),
+                        key=checkbox_key
+                    )
+                with action_cols[1]:
+                    mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.md': 'text/markdown'}
+                    mime_type = mime_map.get(file_ext, "application/octet-stream")
+                    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"Download Error: {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'][gallery_type].pop(file, None)
+                            if file in st.session_state.get('layout_snapshots', []):
+                                st.session_state['layout_snapshots'].remove(file)
+                            logger.info(f"Deleted {gallery_type} asset: {file}")
+                            st.toast(f"Deleted {basename}!", icon="✅")
+                            st.rerun()
+                        except OSError as e:
+                            logger.error(f"Error deleting file {file}: {e}")
+                            st.error(f"Could not delete {basename}")
+            except Exception as e:
+                st.error(f"Error displaying {basename}: {e}")
+                logger.error(f"Error displaying asset {file}: {e}")
+            st.markdown("---")
+
+# --- UI Elements ---
+st.sidebar.subheader("🤖 AI Settings")
+with st.sidebar.expander("API Inference Settings", expanded=False):
+    st.session_state.hf_custom_key = st.text_input(
+        "Custom HF Token",
+        value=st.session_state.get('hf_custom_key', ""),
+        type="password",
+        key="hf_custom_key_input"
+    )
+    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"HF Token Status: {token_status}")
+    providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
+    st.session_state.hf_provider = st.selectbox(
+        "HF Inference Provider",
+        options=providers_list,
+        index=providers_list.index(st.session_state.get('hf_provider', DEFAULT_PROVIDER)),
+        key="hf_provider_select"
+    )
+    st.session_state.hf_custom_api_model = st.text_input(
+        "Custom HF API Model ID",
+        value=st.session_state.get('hf_custom_api_model', ""),
+        key="hf_custom_model_input"
+    )
+    effective_hf_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 HF 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"
+    )
+    st.caption(f"Effective HF API Model: {effective_hf_model}")
+    if _openai_available:
+        st.session_state.openai_selected_model = st.selectbox(
+            "OpenAI Model",
+            options=OPENAI_MODELS_LIST,
+            index=OPENAI_MODELS_LIST.index(st.session_state.get('openai_selected_model', OPENAI_MODELS_LIST[0])),
+            key="openai_model_select"
+        )
+
+with st.sidebar.expander("Local Model Selection", expanded=True):
+    if not _transformers_available:
+        st.warning("Transformers library not found. Cannot load 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', "None")
+        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"
+        )
+        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.")
+
+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), 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")
+    st.session_state.gen_frequency_penalty = st.slider("Repetition Penalty", 0.0, 1.0, st.session_state.get('gen_frequency_penalty', 0.0), step=0.05, key="param_repetition")
+    st.session_state.gen_seed = st.slider("Seed", -1, 65535, st.session_state.get('gen_seed', -1), step=1, key="param_seed")
+
+st.sidebar.subheader("🖼️ Gallery Settings")
+st.slider(
+    "Max Items Shown",
+    min_value=2,
+    max_value=50,
+    value=st.session_state.get('gallery_size', 10),
+    key="gallery_size_slider"
+)
+st.session_state.gallery_size = st.session_state.gallery_size_slider
+st.sidebar.markdown("---")
+update_gallery('image')
+update_gallery('md')
+update_gallery('pdf')
+
+# --- Main Application ---
+st.title("Vision & Layout Titans 🚀🖼️📄")
+st.markdown("Create PDFs from images and markdown, process with AI, and manage characters.")
+tabs = st.tabs([
+    "Image/MD->PDF Layout 🖼️➡️📄",
+    "Camera Snap 📷",
+    "Download PDFs 📥",
+    "Build Titan (Local Models) 🌱",
+    "PDF Process (AI) 📄",
+    "Image Process (AI) 🖼️",
+    "Text Process (AI) 📝",
+    "Test OCR (AI) 🔍",
+    "Test Image Gen (Diffusers) 🎨",
+    "Character Editor 🧑‍🎨",
+    "Character Gallery 🖼️"
+])
+
+with tabs[0]:
+    st.header("Image/Markdown to PDF Layout Generator")
+    st.markdown("Select images and markdown files, reorder them, and generate a PDF.")
+    col1, col2 = st.columns(2)
+    with col1:
+        st.subheader("A. Select Assets")
+        selected_images = [f for f in get_typed_gallery_files('image') if st.session_state['asset_checkboxes']['image'].get(f, False)]
+        selected_mds = [f for f in get_typed_gallery_files('md') if st.session_state['asset_checkboxes']['md'].get(f, False)]
+        st.write(f"Selected Images: {len(selected_images)}")
+        st.write(f"Selected Markdown Files: {len(selected_mds)}")
+    with col2:
+        st.subheader("B. Review and Reorder")
+        layout_records = []
+        for idx, path in enumerate(selected_images + selected_mds, start=1):
+            is_md = path in selected_mds
+            try:
+                if is_md:
+                    with open(path, 'r', encoding='utf-8') as f:
+                        content = f.read(50)
+                    layout_records.append({
+                        "filename": os.path.basename(path),
+                        "source": path,
+                        "type": "Markdown",
+                        "preview": content + "...",
+                        "order": idx
+                    })
+                else:
+                    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,
+                        "type": "Image",
+                        "width": w,
+                        "height": h,
+                        "aspect_ratio": ar,
+                        "orientation": orient,
+                        "order": idx
+                    })
+            except Exception as e:
+                logger.warning(f"Could not process {path}: {e}")
+                st.warning(f"Skipping invalid file: {os.path.basename(path)}")
+        if not layout_records:
+            st.infoperiod