diff --git "a/app.py" "b/app.py" --- "a/app.py" +++ "b/app.py" @@ -1,4 +1,6 @@ -# --- Combined Imports ------------------------------------ +#!/usr/bin/env python +# app.py + import io import os import re @@ -11,1683 +13,276 @@ import time import zipfile import json import asyncio -import aiofiles +from pathlib import Path from datetime import datetime -from collections import Counter -from dataclasses import dataclass, field -from io import BytesIO -from typing import Optional, List, Dict, Any +from typing import Any, List, Dict, Optional import pandas as pd import pytz import streamlit as st -from PIL import Image, ImageDraw # Added ImageDraw +import aiofiles +import requests + +from PIL import Image, ImageDraw, UnidentifiedImageError 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 🌌" - } -) +from reportlab.lib.pagesizes import letter +import fitz # PyMuPDF +from huggingface_hub import InferenceClient +from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError -# Conditional imports for optional/heavy libraries +# Optional AI/ML imports try: import torch - from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline - # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.) + 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 - # Don't show warning if transformers also missing, handled above - if _transformers_available: - st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.") +# --- Page 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', + 'About': "Combined App: Imageβ†’PDF Layout + HF AI Tools 🌌" + } +) -import requests # Keep requests import - -# --- Logging Setup --------------------------------------- +# --- Logging Setup --- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) -log_records = [] +log_records: List[logging.LogRecord] = [] class LogCaptureHandler(logging.Handler): def emit(self, record): log_records.append(record) logger.addHandler(LogCaptureHandler()) -# --- Environment Variables & Constants ------------------- +# --- Constants & Defaults --- 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 + "meta-llama/Meta-Llama-3.1-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", - "google/gemma-2-9b-it", # Added Gemma 2 - "Qwen/Qwen2-7B-Instruct", # Added Qwen2 + "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", # Larger Mixture of Experts - # Add a smaller option + "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", "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() +# --- Session State Initialization --- +def _init_state(key: str, default: Any): + if key not in st.session_state: + st.session_state[key] = default + +for k, v in { + 'layout_snapshots': [], + 'layout_new_uploads': [], + 'layout_last_capture': None, + 'history': [], + 'processing': {}, + 'asset_checkboxes': {}, + 'downloaded_pdfs': {}, + 'unique_counter': 0, + 'cam0_file': None, + 'cam1_file': None, + 'characters': [], + 'char_form_reset_key': 0, + 'gallery_size': 10, + 'hf_inference_client': None, + 'hf_provider': DEFAULT_PROVIDER, + 'hf_custom_key': "", + 'hf_selected_api_model': FEATURED_MODELS_LIST[0], + 'hf_custom_api_model': "", + 'local_models': {}, + 'selected_local_model_path': None, + 'gen_max_tokens': 512, + 'gen_temperature': 0.7, + 'gen_top_p': 0.95, + 'gen_frequency_penalty': 0.0, + 'gen_seed': -1 +}.items(): + _init_state(k, v) + +# --- Utility Functions --- +def generate_filename(seq: str, ext: str = "png") -> str: + ts = time.strftime('%Y%m%d_%H%M%S') + safe = re.sub(r'[^\w\-]+', '_', seq) + return f"{safe}_{ts}.{ext}" -# --- 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'{label}' - 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): +def clean_stem(fn: str) -> str: + return os.path.splitext(os.path.basename(fn))[0].replace('-', ' ').replace('_', ' ').title() + +def get_download_link(path: str, mime: str, label: str = "Download") -> str: + if not os.path.exists(path): return f"{label} (not found)" + data = open(path,'rb').read() + b64 = base64.b64encode(data).decode() + return f'{label}' + +def get_gallery_files(types: List[str] = ['png','jpg','jpeg','pdf','md','txt']) -> List[str]: + files = set() + for ext in types: + files.update(glob.glob(f"*.{ext}")) + files.update(glob.glob(f"*.{ext.upper()}")) + return sorted(files) + +# Delete with rerun +def delete_asset(path: str): 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 = [] + os.remove(path) + st.session_state['asset_checkboxes'].pop(path, None) + if path in st.session_state['layout_snapshots']: + st.session_state['layout_snapshots'].remove(path) + st.toast(f"Deleted {os.path.basename(path)}", icon="βœ…") + except OSError as e: + st.error(f"Delete failed: {e}") + st.rerun() + +# Sidebar gallery updater +def update_gallery(): + st.sidebar.markdown("### Asset Gallery πŸ“ΈπŸ“–") + files = get_gallery_files() + if not files: + st.sidebar.info("No assets.") + return + st.sidebar.caption(f"Found {len(files)} assets.") + for f in files[:st.session_state['gallery_size']]: + name = os.path.basename(f) + ext = os.path.splitext(f)[1].lower() + st.sidebar.markdown(f"**{name}**") + with st.sidebar.expander("Preview", expanded=False): + try: + if ext in ['.png','.jpg','.jpeg']: + st.image(Image.open(f), use_container_width=True) + elif ext == '.pdf': + doc = fitz.open(f) + if doc.page_count: + 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) + doc.close() + else: + txt = Path(f).read_text(errors='ignore') + st.code(txt[:200]+'…') + except: + st.warning("Preview error") + c1,c2,c3 = st.sidebar.columns(3) + sel = st.session_state['asset_checkboxes'].get(f, False) + c1.checkbox("Select", value=sel, key=f"cb_{f}") + st.session_state['asset_checkboxes'][f] = st.session_state.get(f"cb_{f}") + mime = {'png':'image/png','jpg':'image/jpeg','jpeg':'image/jpeg','pdf':'application/pdf','md':'text/markdown','txt':'text/plain'}.get(ext[1:], 'application/octet-stream') + with open(f,'rb') as fp: + c2.download_button("πŸ“₯", data=fp, file_name=name, mime=mime, key=f"dl_{f}") + c3.button("πŸ—‘οΈ", key=f"del_{f}", on_click=delete_asset, args=(f,)) + st.sidebar.markdown("---") + +# --- PDF Snapshot & Generation --- +async def process_pdf_snapshot(path: str, mode: str='single', resF: float=2.0) -> List[str]: + status = st.empty() + status.text("Snapshot start...") + out_files: List[str] = [] 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() + doc = fitz.open(path) + mat = fitz.Matrix(resF,resF) + cnt = {'single':1,'twopage':2,'allpages':len(doc)}.get(mode,1) + for i in range(min(cnt,len(doc))): + s = 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) - + pix = page.get_pixmap(matrix=mat) + base = os.path.splitext(os.path.basename(path))[0] + fname = generate_filename(f"{base}_pg{i+1}_{mode}","png") + await asyncio.to_thread(pix.save, fname) + out_files.append(fname) + status.text(f"Saved {fname} ({int(time.time()-s)}s)") 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 + status.success(f"Snapshot done: {len(out_files)} 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: + status.error(f"Snapshot error: {e}") + for f in out_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).") + out_files = [] + return out_files + +from reportlab.lib.pagesizes import letter + +def make_image_sized_pdf(sources: List[Any]) -> Optional[bytes]: + # dedupe + seen, uniq = set(), [] + for s in sources: + key = s if isinstance(s,str) else getattr(s,'name',None) + if key and key not in seen: + seen.add(key) + uniq.append(s) + if not uniq: + st.warning("No images for PDF") 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: + buf = io.BytesIO() + c = canvas.Canvas(buf, pagesize=letter) + status = st.empty() + for idx,s in enumerate(uniq,1): 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.") + img = Image.open(s) if isinstance(s,str) else Image.open(s) + w,h = img.size + cap = 30 + c.setPageSize((w,h+cap)) + c.drawImage(ImageReader(img),0,cap,w,h,mask='auto') + cap_txt = clean_stem(s if isinstance(s,str) else s.name) + c.setFont('Helvetica',12) + c.drawCentredString(w/2,cap/2,cap_txt) + c.setFont('Helvetica',8) + c.drawRightString(w-10,10,str(idx)) + c.showPage() + status.text(f"Page {idx}/{len(uniq)} added") 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 + status.error(f"Error page {idx}: {e}") + c.save() + buf.seek(0) + return buf.getvalue() - return st.session_state.hf_inference_client - -# --- HF/Local Model Processing Functions (Replaced OpenAI ones) --- +# --- HF Inference Client --- +def get_hf_client() -> Optional[InferenceClient]: + provider = st.session_state['hf_provider'] + token = st.session_state['hf_custom_key'].strip() or HF_TOKEN + if provider!='hf-inference' and not token: + st.error(f"Provider {provider} needs token") + return None + client = st.session_state['hf_inference_client'] + if not client: + st.session_state['hf_inference_client'] = InferenceClient(token=token, provider=provider) + return st.session_state['hf_inference_client'] +# --- HF Processing --- 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]" - + stp = st.empty(); stp.text("Processing...") + msgs = [{"role":"system","content":"You are an assistant."}, + {"role":"user","content":f"{prompt}\n\n{text}"}] + out = "" 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 - + if not client: return "Client error" + model = st.session_state['hf_custom_api_model'] or st.session_state['hf_selected_api_model'] 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") \ No newline at end of file + resp = client.chat_completion( + model=model, + messages=msgs, + max_tokens=st.session_state['gen_max_tokens'], + temperature=st.session +]}]}