diff --git "a/app.py" "b/app.py" --- "a/app.py" +++ "b/app.py" @@ -15,59 +15,46 @@ import aiofiles from datetime import datetime from collections import Counter -from dataclasses import dataclass +from dataclasses import dataclass, field from io import BytesIO -from typing import Optional +from typing import Optional, List, Dict, Any import pandas as pd import pytz import streamlit as st -from PIL import Image +from PIL import Image, ImageDraw # Added ImageDraw from reportlab.pdfgen import canvas from reportlab.lib.utils import ImageReader +from reportlab.lib.pagesizes import letter # Default page size import fitz # PyMuPDF -# --- App Configuration (Choose one, adapted from App 2) --- -st.set_page_config( - page_title="Vision & Layout Titans ๐Ÿš€๐Ÿ–ผ๏ธ", - page_icon="๐Ÿค–", - layout="wide", - initial_sidebar_state="expanded", - menu_items={ - 'Get Help': 'https://huggingface.co/awacke1', - 'Report a Bug': 'https://huggingface.co/spaces/awacke1', - 'About': "Combined App: Image->PDF Layout + AI Vision & SFT Titans ๐ŸŒŒ" - } -) +# --- Hugging Face Imports --- +from huggingface_hub import InferenceClient, HfApi, list_models +from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Import specific exceptions # Conditional imports for optional/heavy libraries try: import torch - from diffusers import StableDiffusionPipeline - from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel - _ai_libs_available = True + from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline + # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.) + _transformers_available = True except ImportError: - _ai_libs_available = False - st.sidebar.warning("AI/ML libraries (torch, transformers, diffusers) not found. Some AI features disabled.") + _transformers_available = False + st.sidebar.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.") try: - from openai import OpenAI - client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'), organization=os.getenv('OPENAI_ORG_ID')) - _openai_available = True - if not os.getenv('OPENAI_API_KEY'): - st.sidebar.warning("OpenAI API Key/Org ID not found in environment variables. GPT features disabled.") - _openai_available = False + from diffusers import StableDiffusionPipeline + _diffusers_available = True except ImportError: - _openai_available = False - st.sidebar.warning("OpenAI library not found. GPT features disabled.") -except Exception as e: - _openai_available = False - st.sidebar.warning(f"OpenAI client error: {e}. GPT features disabled.") + _diffusers_available = False + # Don't show warning if transformers also missing, handled above + if _transformers_available: + st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.") import requests # Keep requests import -# --- Logging Setup (from App 2) -------------------------- +# --- Logging Setup --------------------------------------- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) log_records = [] @@ -76,42 +63,108 @@ class LogCaptureHandler(logging.Handler): log_records.append(record) logger.addHandler(LogCaptureHandler()) -# --- Session State Initialization (Combined) ------------- -# From App 1 -st.session_state.setdefault('layout_snapshots', []) # Renamed to avoid potential conflict +# --- Environment Variables & Constants ------------------- +HF_TOKEN = os.getenv("HF_TOKEN") +DEFAULT_PROVIDER = "hf-inference" +# Model List (curated, similar to Gradio example) - can be updated +FEATURED_MODELS_LIST = [ + "meta-llama/Meta-Llama-3.1-8B-Instruct", # Updated Llama model + "mistralai/Mistral-7B-Instruct-v0.3", + "google/gemma-2-9b-it", # Added Gemma 2 + "Qwen/Qwen2-7B-Instruct", # Added Qwen2 + "microsoft/Phi-3-mini-4k-instruct", + "HuggingFaceH4/zephyr-7b-beta", + "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", # Larger Mixture of Experts + # Add a smaller option + "HuggingFaceTB/SmolLM-1.7B-Instruct" +] +# Add common vision models if planning local loading +VISION_MODELS_LIST = [ + "Salesforce/blip-image-captioning-large", + "microsoft/trocr-large-handwritten", # OCR model + "llava-hf/llava-1.5-7b-hf", # Vision Language Model + "google/vit-base-patch16-224", # Basic Vision Transformer +] +DIFFUSION_MODELS_LIST = [ + "stabilityai/stable-diffusion-xl-base-1.0", # Common SDXL + "runwayml/stable-diffusion-v1-5", # Classic SD 1.5 + "OFA-Sys/small-stable-diffusion-v0", # Tiny diffusion +] + + +# --- 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 ๐ŸŒŒ" + } +) + +# --- Session State Initialization (Combined & Updated) --- +# Layout PDF specific +st.session_state.setdefault('layout_snapshots', []) +st.session_state.setdefault('layout_new_uploads', []) -# From App 2 +# General App State st.session_state.setdefault('history', []) -st.session_state.setdefault('builder', None) -st.session_state.setdefault('model_loaded', False) 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('selected_model_type', "Causal LM") -st.session_state.setdefault('selected_model', "None") 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', False) +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() -st.session_state.setdefault('gallery_size', 2) # From App 2 gallery settings -# --- Dataclasses (from App 2) ---------------------------- +# --- Dataclasses (Refined for Local Models) ------------- @dataclass -class ModelConfig: - name: str - base_model: str - size: str +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 - model_type: str = "causal_lm" - @property - def model_path(self): - return f"models/{self.name}" + 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: +class DiffusionConfig: # Kept for clarity in diffusion tab if needed name: str base_model: str size: str @@ -120,97 +173,24 @@ class DiffusionConfig: def model_path(self): return f"diffusion_models/{self.name}" -# --- Class Definitions (from App 2) ----------------------- -# Simplified ModelBuilder and DiffusionBuilder if libraries are missing -if _ai_libs_available: - class ModelBuilder: - def __init__(self): - self.config = None - self.model = None - self.tokenizer = None - self.jokes = [ - "Why did the AI go to therapy? Too many layers to unpack! ๐Ÿ˜‚", - "Training complete! Time for a binary coffee break. โ˜•", - "I told my neural network a joke; it couldn't stop dropping bits! ๐Ÿค–", - "I asked the AI for a pun, and it said, 'I'm punning on parallel processing!' ๐Ÿ˜„", - "Debugging my code is like a stand-up routineโ€”always a series of exceptions! ๐Ÿ˜†" - ] - def load_model(self, model_path: str, config: Optional[ModelConfig] = None): - with st.spinner(f"Loading {model_path}... โณ"): - self.model = AutoModelForCausalLM.from_pretrained(model_path) - self.tokenizer = AutoTokenizer.from_pretrained(model_path) - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - if config: - self.config = config - self.model.to("cuda" if torch.cuda.is_available() else "cpu") - st.success(f"Model loaded! ๐ŸŽ‰ {random.choice(self.jokes)}") - return self - def save_model(self, path: str): - with st.spinner("Saving model... ๐Ÿ’พ"): - os.makedirs(os.path.dirname(path), exist_ok=True) - self.model.save_pretrained(path) - self.tokenizer.save_pretrained(path) - st.success(f"Model saved at {path}! โœ…") - - class DiffusionBuilder: - def __init__(self): - self.config = None - self.pipeline = None - def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None): - with st.spinner(f"Loading diffusion model {model_path}... โณ"): - # Use float32 for broader compatibility, esp. CPU - self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cuda" if torch.cuda.is_available() else "cpu") - if config: - self.config = config - st.success("Diffusion model loaded! ๐ŸŽจ") - return self - def save_model(self, path: str): - with st.spinner("Saving diffusion model... ๐Ÿ’พ"): - os.makedirs(os.path.dirname(path), exist_ok=True) - self.pipeline.save_pretrained(path) - st.success(f"Diffusion model saved at {path}! โœ…") - def generate(self, prompt: str): - # Adjust steps for CPU if needed - steps = 10 if torch.cuda.is_available() else 5 # Fewer steps for CPU demo - with st.spinner(f"Generating image with {steps} steps..."): - image = self.pipeline(prompt, num_inference_steps=steps).images[0] - return image -else: # Placeholder classes if AI libs are missing - class ModelBuilder: - def __init__(self): st.error("AI Libraries not available.") - def load_model(self, *args, **kwargs): pass - def save_model(self, *args, **kwargs): pass - class DiffusionBuilder: - def __init__(self): st.error("AI Libraries not available.") - def load_model(self, *args, **kwargs): pass - def save_model(self, *args, **kwargs): pass - def generate(self, *args, **kwargs): return Image.new("RGB", (64,64), "gray") - # --- 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"): - # Use App 2's more robust version timestamp = time.strftime('%Y%m%d_%H%M%S') - # Sanitize sequence name for filename safe_sequence = re.sub(r'[^\w\-]+', '_', str(sequence)) return f"{safe_sequence}_{timestamp}.{ext}" def pdf_url_to_filename(url): - # Use App 2's version - # Further sanitize - remove http(s) prefix and limit length 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"): - # Use App 2's version, ensure file exists - if not os.path.exists(file_path): - return f"{label} (File not found)" + 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() + with open(file_path, "rb") as f: file_bytes = f.read() b64 = base64.b64encode(file_bytes).decode() return f'{label}' except Exception as e: @@ -218,64 +198,49 @@ def get_download_link(file_path, mime_type="application/octet-stream", label="Do return f"{label} (Error)" def zip_directory(directory_path, zip_path): - # Use App 2's version 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_model_files(model_type="causal_lm"): - # Use App 2's version - pattern = "models/*" if model_type == "causal_lm" else "diffusion_models/*" +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 if dirs else ["None"] + return dirs -def get_gallery_files(file_types=("png", "pdf", "jpg", "jpeg", "md", "txt")): # Expanded types - # Use App 2's version, ensure lowercase extensions +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()}")) # Include uppercase extensions too + all_files.update(glob.glob(f"*.{ext.upper()}")) return sorted(list(all_files)) def get_pdf_files(): - # Use App 2's version return sorted(glob.glob("*.pdf") + glob.glob("*.PDF")) def download_pdf(url, output_path): - # Use App 2's version try: - headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'} - response = requests.get(url, stream=True, timeout=20, headers=headers) # Added user-agent, longer timeout - response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx) + 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) + 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}") - # Attempt to remove partially downloaded file - if os.path.exists(output_path): - try: - os.remove(output_path) - logger.info(f"Removed partially downloaded file: {output_path}") - except OSError as remove_error: - logger.error(f"Error removing partial file {output_path}: {remove_error}") + 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 + 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): - # Use App 2's version, added resolution control start_time = time.time() - # Use a placeholder within the main app area for status during async operations status_placeholder = st.empty() status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... (0s)") output_files = [] @@ -283,25 +248,21 @@ async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): 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) + if mode == "single": num_pages_to_process = min(1, len(doc)) + elif mode == "twopage": num_pages_to_process = min(2, len(doc)) + elif mode == "allpages": num_pages_to_process = len(doc) for i in range(num_pages_to_process): page_start_time = time.time() page = doc[i] pix = page.get_pixmap(matrix=matrix) - # Use PDF name and page number in filename for clarity 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) # Run sync save in thread + 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) # Yield control briefly + await asyncio.sleep(0.01) doc.close() elapsed = int(time.time() - start_time) @@ -310,164 +271,310 @@ async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): 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}") - # Clean up any files created before the error for f in output_files: if os.path.exists(f): os.remove(f) return [] -async def process_gpt4o_ocr(image: Image.Image, output_file: str): - # Use App 2's version, check for OpenAI availability - if not _openai_available: - st.error("OpenAI OCR requires API key and library.") - return "" - start_time = time.time() - status_placeholder = st.empty() - status_placeholder.text("Processing GPT-4o OCR... (0s)") - buffered = BytesIO() - # Ensure image is in a compatible format (e.g., PNG, JPEG) - save_format = "PNG" if image.format != "JPEG" else "JPEG" - image.save(buffered, format=save_format) - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - messages = [{ - "role": "user", - "content": [ - {"type": "text", "text": "Extract text content from the image. Provide only the extracted text."}, # More specific prompt - {"type": "image_url", "image_url": {"url": f"data:image/{save_format.lower()};base64,{img_str}", "detail": "auto"}} - ] - }] - try: - # Run OpenAI call in a separate thread to avoid blocking Streamlit's event loop - response = await asyncio.to_thread( - client.chat.completions.create, - model="gpt-4o", messages=messages, max_tokens=4000 # Increased tokens - ) - result = response.choices[0].message.content or "" # Handle potential None result - elapsed = int(time.time() - start_time) - status_placeholder.success(f"GPT-4o OCR completed in {elapsed}s!") - async with aiofiles.open(output_file, "w", encoding='utf-8') as f: # Specify encoding - await f.write(result) - logger.info(f"GPT-4o OCR successful for {output_file}") - return result - except Exception as e: - logger.error(f"Failed to process image with GPT-4o: {e}") - status_placeholder.error(f"GPT-4o OCR Failed: {e}") - return f"Error during OCR: {str(e)}" +# --- HF Inference Client Management --- +def get_hf_client() -> Optional[InferenceClient]: + """Gets or initializes the Hugging Face Inference Client based on session state.""" + provider = st.session_state.hf_provider + custom_key = st.session_state.hf_custom_key.strip() + token_to_use = custom_key if custom_key else HF_TOKEN + if not token_to_use and provider != "hf-inference": + st.error(f"Provider '{provider}' requires a Hugging Face API token (either via HF_TOKEN env var or custom key).") + return None + if provider == "hf-inference" and not token_to_use: + logger.warning("Using hf-inference provider without a token. Rate limits may apply.") + token_to_use = None # Explicitly set to None for public inference API + + # Check if client needs re-initialization + current_client = st.session_state.get('hf_inference_client') + # Simple check: re-init if provider or token presence changes + needs_reinit = True + if current_client: + # Basic check, more robust checks could compare client._token etc. if needed + # This assumes provider and token status are the key determinants + client_uses_custom = hasattr(current_client, '_token') and current_client._token == custom_key + client_uses_default = hasattr(current_client, '_token') and current_client._token == HF_TOKEN + client_uses_no_token = not hasattr(current_client, '_token') or current_client._token is None + + if current_client.provider == provider: + if custom_key and client_uses_custom: needs_reinit = False + elif not custom_key and HF_TOKEN and client_uses_default: needs_reinit = False + elif not custom_key and not HF_TOKEN and client_uses_no_token: needs_reinit = False + + + if needs_reinit: + try: + logger.info(f"Initializing InferenceClient for provider: {provider}. Token source: {'Custom Key' if custom_key else ('HF_TOKEN' if HF_TOKEN else 'None')}") + st.session_state.hf_inference_client = InferenceClient(token=token_to_use, provider=provider) + logger.info("InferenceClient initialized successfully.") + except Exception as e: + st.error(f"Failed to initialize Hugging Face client for provider {provider}: {e}") + logger.error(f"InferenceClient initialization failed: {e}") + st.session_state.hf_inference_client = None -async def process_image_gen(prompt: str, output_file: str): - # Use App 2's version, check AI lib availability - if not _ai_libs_available: - st.error("Image Generation requires AI libraries.") - img = Image.new("RGB", (256, 256), "lightgray") - draw = ImageDraw.Draw(img) - draw.text((10, 10), "AI libs missing", fill="black") - img.save(output_file) - return img + return st.session_state.hf_inference_client - start_time = time.time() +# --- HF/Local Model Processing Functions (Replaced OpenAI ones) --- + +def process_text_hf(text: str, prompt: str, use_api: bool) -> str: + """Processes text using either HF Inference API or a loaded local model.""" status_placeholder = st.empty() - status_placeholder.text("Processing Image Gen... (0s)") + 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)}" - # Ensure a pipeline is loaded, default to small one if necessary - pipeline = None - if st.session_state.get('builder') and isinstance(st.session_state['builder'], DiffusionBuilder) and st.session_state['builder'].pipeline: - pipeline = st.session_state['builder'].pipeline 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: - with st.spinner("Loading default small diffusion model..."): - pipeline = StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cuda" if torch.cuda.is_available() else "cpu") - st.info("Loaded default small diffusion model for image generation.") + # 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"Failed to load default diffusion model: {e}") - status_placeholder.error(f"Failed to load default diffusion model: {e}") - img = Image.new("RGB", (256, 256), "lightgray") - draw = ImageDraw.Draw(img) - draw.text((10, 10), "Model load error", fill="black") - img.save(output_file) - return img + logger.error(f"Local text processing failed for model {model_path}: {e}") + result_text = f"Error during local model inference: {str(e)}" - try: - # Run generation in a thread - gen_image = await asyncio.to_thread(pipeline, prompt, num_inference_steps=15) # Slightly more steps - gen_image = gen_image.images[0] # Extract image from list - elapsed = int(time.time() - start_time) - status_placeholder.success(f"Image Gen completed in {elapsed}s!") - await asyncio.to_thread(gen_image.save, output_file) # Save in thread - logger.info(f"Image generation successful for {output_file}") - return gen_image - except Exception as e: - logger.error(f"Image generation failed: {e}") - status_placeholder.error(f"Image generation failed: {e}") - # Create placeholder error image - img = Image.new("RGB", (256, 256), "lightgray") - from PIL import ImageDraw - draw = ImageDraw.Draw(img) - draw.text((10, 10), f"Generation Error:\n{e}", fill="red") - await asyncio.to_thread(img.save, output_file) - return img - -# --- GPT Processing Functions (from App 2, with checks) --- -def process_image_with_prompt(image: Image.Image, prompt: str, model="gpt-4o-mini", detail="auto"): - if not _openai_available: return "Error: OpenAI features disabled." - status_placeholder = st.empty() - status_placeholder.info(f"Processing image with GPT ({model})...") - buffered = BytesIO() - save_format = "PNG" if image.format != "JPEG" else "JPEG" - image.save(buffered, format=save_format) - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - messages = [{ - "role": "user", - "content": [ - {"type": "text", "text": prompt}, - {"type": "image_url", "image_url": {"url": f"data:image/{save_format.lower()};base64,{img_str}", "detail": detail}} - ] - }] - try: - response = client.chat.completions.create(model=model, messages=messages, max_tokens=1000) # Increased tokens - result = response.choices[0].message.content or "" - status_placeholder.success(f"GPT ({model}) image processing complete.") - logger.info(f"GPT ({model}) image processing successful.") - return result - except Exception as e: - logger.error(f"Error processing image with GPT ({model}): {e}") - status_placeholder.error(f"Error processing image with GPT ({model}): {e}") - return f"Error processing image with GPT: {str(e)}" + elapsed = int(time.time() - start_time) + status_placeholder.success(f"Text processing completed in {elapsed}s.") + return result_text + -def process_text_with_prompt(text: str, prompt: str, model="gpt-4o-mini"): - if not _openai_available: return "Error: OpenAI features disabled." +# --- 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() - status_placeholder.info(f"Processing text with GPT ({model})...") - messages = [{"role": "user", "content": f"{prompt}\n\n---\n\n{text}"}] # Added separator - try: - response = client.chat.completions.create(model=model, messages=messages, max_tokens=2000) # Increased tokens - result = response.choices[0].message.content or "" - status_placeholder.success(f"GPT ({model}) text processing complete.") - logger.info(f"GPT ({model}) text processing successful.") - return result - except Exception as e: - logger.error(f"Error processing text with GPT ({model}): {e}") - status_placeholder.error(f"Error processing text with GPT ({model}): {e}") - return f"Error processing text with GPT: {str(e)}" + start_time = time.time() + result_text = "[Image processing not fully implemented with HF models yet]" + + if use_api: + # --- Use HF API (Basic Image-to-Text Example) --- + status_placeholder.info("Processing image using Hugging Face API (Image-to-Text)...") + client = get_hf_client() + if not client: return "Error: HF client not configured." + + # Convert PIL image to bytes + buffered = BytesIO() + image.save(buffered, format="PNG" if image.format != 'JPEG' else 'JPEG') + img_bytes = buffered.getvalue() + + try: + # Example using a generic image-to-text model via API + # NOTE: This does NOT use the 'prompt' effectively like VQA models. + # Need to select an appropriate model ID known for image captioning. + # Using a default BLIP model for demonstration. + captioning_model_id = "Salesforce/blip-image-captioning-large" + status_placeholder.info(f"Using API Image-to-Text Model: {captioning_model_id}") + + response_list = client.image_to_text(data=img_bytes, model=captioning_model_id) + + if response_list and isinstance(response_list, list) and 'generated_text' in response_list[0]: + result_text = f"API Caption ({captioning_model_id}): {response_list[0]['generated_text']}\n\n(Note: API call did not use custom prompt: '{prompt}')" + logger.info(f"HF API image captioning successful for model {captioning_model_id}.") + else: + result_text = "Error: Unexpected response format from image-to-text API." + logger.warning(f"Unexpected API response for image-to-text: {response_list}") + + except Exception as e: + logger.error(f"HF API image processing failed: {e}") + result_text = f"Error during Hugging Face API image inference: {str(e)}" + + else: + # --- Use Local Vision Model --- + status_placeholder.info("Processing image using local model...") + if not _transformers_available: return "Error: Transformers library needed." + + model_path = st.session_state.get('selected_local_model_path') + if not model_path or model_path not in st.session_state.get('local_models', {}): + return "Error: No suitable local model selected or loaded." + + local_model_data = st.session_state['local_models'][model_path] + model_type = local_model_data.get('type') + + # --- Placeholder Logic - Requires Specific Model Implementation --- + if model_type == 'vision': # General VQA or Captioning + status_placeholder.warning(f"Local Vision Model ({os.path.basename(model_path)}): Processing logic depends heavily on the specific model architecture (e.g., LLaVA, BLIP). Placeholder implementation.") + # Example: Needs processor + model.generate based on model type + # processor = local_model_data.get('processor') + # model = local_model_data.get('model') + # if processor and model: + # try: + # # inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device) + # # generated_ids = model.generate(**inputs, max_new_tokens=...) + # # result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() + # result_text = f"[Local vision processing for {os.path.basename(model_path)} needs specific implementation based on its type.] Prompt was: {prompt}" + # except Exception as e: + # result_text = f"Error during local vision model inference: {e}" + # else: + # result_text = "Error: Processor or model missing for local vision task." + result_text = f"[Local vision processing for {os.path.basename(model_path)} needs specific implementation based on its type.] Prompt was: {prompt}" # Placeholder + + elif model_type == 'ocr': # OCR Specific Model + status_placeholder.warning(f"Local OCR Model ({os.path.basename(model_path)}): Placeholder implementation.") + # Example for TrOCR style models + # processor = local_model_data.get('processor') + # model = local_model_data.get('model') + # if processor and model: + # try: + # # pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(model.device) + # # generated_ids = model.generate(pixel_values, max_new_tokens=...) + # # result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + # result_text = f"[Local OCR processing for {os.path.basename(model_path)} needs specific implementation.]" + # except Exception as e: + # result_text = f"Error during local OCR model inference: {e}" + # else: + # result_text = "Error: Processor or model missing for local OCR task." + result_text = f"[Local OCR processing for {os.path.basename(model_path)} needs specific implementation.]" # Placeholder + else: + result_text = f"Error: Loaded model '{os.path.basename(model_path)}' is not a recognized vision/OCR type for this function." + + + elapsed = int(time.time() - start_time) + status_placeholder.success(f"Image processing attempt completed in {elapsed}s.") + return result_text + +# Basic OCR function using the image processor above +async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool) -> str: + """ Performs OCR using the process_image_hf function framework. """ + # Simple prompt for OCR task + ocr_prompt = "Extract text content from this image." + result = process_image_hf(image, ocr_prompt, use_api) -# --- Character Functions (from App 2) -------------------- + # 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(): - # Use App 2's version intro_templates = [ - "{char} is a valiant knight who is silent and reserved, he looks handsome but aloof.", - "{char} is a mischievous thief with a heart of gold, always sneaking around but helping those in need.", - "{char} is a wise scholar who loves books more than people, often lost in thought.", - "{char} is a fiery warrior with a short temper, but fiercely loyal to friends.", - "{char} is a gentle healer who speaks softly, always carrying herbs and a warm smile." + "{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 = [ - "You were startled by the sudden intrusion of a man into your home. 'I am from the knight's guild, and I have been ordered to arrest you.'", - "A shadowy figure steps into the light. 'I heard you needed helpโ€”nameโ€™s {char}, best thief in town.'", - "A voice calls from behind a stack of books. 'Oh, hello! Iโ€™m {char}, didnโ€™t see you thereโ€”too many scrolls!'", - "A booming voice echoes, 'Iโ€™m {char}, and Iโ€™m here to fight for justiceโ€”or at least a good brawl!'", - "A soft hand touches your shoulder. 'Iโ€™m {char}, here to heal your woundsโ€”donโ€™t worry, Iโ€™ve got you.'" - ] + "'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) @@ -475,17 +582,14 @@ def randomize_character_content(): return name, gender, intro, greeting def save_character(character_data): - # Use App 2's version characters = st.session_state.get('characters', []) - # Prevent duplicate names 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) # Added indent for readability + 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: @@ -494,116 +598,62 @@ def save_character(character_data): return False def load_characters(): - # Use App 2's version - if not os.path.exists("characters.json"): - st.session_state['characters'] = [] - return + 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) - # Basic validation - 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") # Remove invalid file + 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'] = [] - # Attempt to backup corrupted file 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}") + 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 (from App 1, needed for Image->PDF tab) --- +# --- Utility: Clean stems (Keep from previous) ---------- def clean_stem(fn: str) -> str: - # Make it slightly more robust name = os.path.splitext(os.path.basename(fn))[0] name = name.replace('-', ' ').replace('_', ' ') - # Remove common prefixes/suffixes if desired (optional) - # name = re.sub(r'^(scan|img|image)_?', '', name, flags=re.IGNORECASE) - # name = re.sub(r'_?\d+$', '', name) # Remove trailing numbers - return name.strip().title() # Title case + return name.strip().title() -# --- PDF Creation: Image Sized + Captions (from App 1) --- +# --- 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 + if not sources: st.warning("No image sources provided for PDF generation."); return None buf = io.BytesIO() - # Use A4 size initially, will be overridden per page - c = canvas.Canvas(buf, pagesize=(595.27, 841.89)) # Default A4 + 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: - # Handle both file paths and uploaded file objects - if isinstance(src, str): # path - 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: # uploaded file object (BytesIO wrapper) - src.seek(0) # Ensure reading from start - img_obj = Image.open(src) - filename = getattr(src, 'name', f'uploaded_image_{idx}') - src.seek(0) # Reset again just in case needed later - - with img_obj: # Use context manager for PIL Image + 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 # Increased caption height - # Set page size based on image + caption height - pw, ph = iw, ih + cap_h + 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)) - - # Draw image, ensuring it fits within iw, ih space above caption - # Use ImageReader for efficiency with ReportLab img_reader = ImageReader(img_obj) c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto') - - # Draw Caption (cleaned filename) - caption = clean_stem(filename) - c.setFont('Helvetica', 12) - c.setFillColorRGB(0, 0, 0) # Black text - c.drawCentredString(pw / 2, cap_h / 2 + 3, caption) # Center vertically too - - # Draw Page Number - c.setFont('Helvetica', 8) - c.setFillColorRGB(0.5, 0.5, 0.5) # Gray text - c.drawRightString(pw - 10, 8, f"Page {idx}") - - c.showPage() # Finalize the page + 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: # Check if PDF is basically empty - st.error("PDF generation resulted in an empty file. Check image files.") - return None + 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}") @@ -611,1212 +661,920 @@ def make_image_sized_pdf(sources): return None -# --- Sidebar Gallery Update Function (from App 2) -------- +# --- Sidebar Gallery Update Function (Keep from previous) -------- +# ... (update_gallery function assumed here - no changes needed based on AI backend) ... def update_gallery(): container = st.session_state['asset_gallery_container'] - container.empty() # Clear previous gallery rendering - with container.container(): # Use a container to manage layout + container.empty() + with container.container(): st.markdown("### Asset Gallery ๐Ÿ“ธ๐Ÿ“–") st.session_state['gallery_size'] = st.slider("Max Items Shown", 2, 50, st.session_state.get('gallery_size', 10), key="gallery_size_slider") - cols = st.columns(2) # Use 2 columns in the sidebar - all_files = get_gallery_files() # Get currently available files - - if not all_files: - st.info("No assets (images, PDFs, text files) found yet.") - return - + cols = st.columns(2) + all_files = get_gallery_files() + if not all_files: st.info("No assets found."); return files_to_display = all_files[:st.session_state['gallery_size']] - for idx, file in enumerate(files_to_display): with cols[idx % 2]: - st.session_state['unique_counter'] += 1 - unique_id = st.session_state['unique_counter'] - basename = os.path.basename(file) - st.caption(basename) # Show filename as caption above preview - + st.session_state['unique_counter'] += 1; unique_id = st.session_state['unique_counter']; basename = os.path.basename(file); st.caption(basename) try: file_ext = os.path.splitext(file)[1].lower() - if file_ext in ['.png', '.jpg', '.jpeg']: - st.image(Image.open(file), use_container_width=True) + if file_ext in ['.png', '.jpg', '.jpeg']: st.image(Image.open(file), use_container_width=True) elif file_ext == '.pdf': - doc = fitz.open(file) - # Generate preview only if file opens successfully - 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() + 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 file_ext in ['.md', '.txt']: - 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 - checkbox_key = f"asset_cb_{file}_{unique_id}" - # Use get to safely access potentially missing keys after deletion - st.session_state['asset_checkboxes'][file] = st.checkbox( - "Select", - value=st.session_state['asset_checkboxes'].get(file, False), - key=checkbox_key - ) - - 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") + with open(file, 'r', encoding='utf-8', errors='ignore') as f: content_preview = f.read(200) + st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text') + checkbox_key = f"asset_cb_{file}_{unique_id}"; st.session_state['asset_checkboxes'][file] = st.checkbox("Select", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key) + 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") st.markdown(get_download_link(file, mime_type, "๐Ÿ“ฅ"), unsafe_allow_html=True) - delete_key = f"delete_btn_{file}_{unique_id}" 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.success(f"Deleted {basename}") - st.rerun() # Rerun to refresh the gallery immediately - except OSError as e: - logger.error(f"Error deleting file {file}: {e}") - st.error(f"Could not delete {basename}") - - except (fitz.fitz.FileNotFoundError, FileNotFoundError): - st.error(f"File not found: {basename}") - # Clean up state if file is missing - st.session_state['asset_checkboxes'].pop(file, None) - if file in st.session_state.get('layout_snapshots', []): - st.session_state['layout_snapshots'].remove(file) - - except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: - st.error(f"Corrupt PDF: {basename}") - logger.warning(f"Error opening PDF {file}: {pdf_err}") - except UnidentifiedImageError: - st.error(f"Invalid Image: {basename}") - logger.warning(f"Cannot identify image file {file}") - except Exception as e: - st.error(f"Error: {basename}") - logger.error(f"Error displaying asset {file}: {e}") + os.remove(file); st.session_state['asset_checkboxes'].pop(file, None) + if file in st.session_state.get('layout_snapshots', []): st.session_state['layout_snapshots'].remove(file) + logger.info(f"Deleted asset: {file}"); st.success(f"Deleted {basename}"); st.rerun() + except OSError as e: logger.error(f"Error deleting file {file}: {e}"); st.error(f"Could not delete {basename}") + except (fitz.fitz.FileNotFoundError, FileNotFoundError): st.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.error(f"Corrupt PDF: {basename}"); logger.warning(f"Error opening PDF {file}: {pdf_err}") + except UnidentifiedImageError: st.error(f"Invalid Image: {basename}"); logger.warning(f"Cannot identify image file {file}") + except Exception as e: st.error(f"Error: {basename}"); logger.error(f"Error displaying asset {file}: {e}") + st.markdown("---") + if len(all_files) > st.session_state['gallery_size']: st.caption(f"Showing {st.session_state['gallery_size']} of {len(all_files)} assets.") + + +# --- 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}") - st.markdown("---") # Separator between items - if len(all_files) > st.session_state['gallery_size']: - st.caption(f"Showing {st.session_state['gallery_size']} of {len(all_files)} assets.") +# 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 # --- App Title ------------------------------------------- -st.title("Vision & Layout Titans ๐Ÿš€๐Ÿ–ผ๏ธ๐Ÿ“„") -st.markdown("Combined App: AI Vision/SFT Tools + Image-to-PDF Layout Generator") +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 ๐Ÿ–ผ๏ธโžก๏ธ๐Ÿ“„", # Added from App 1 + "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 ๐Ÿ“ฅ", - "PDF Process ๐Ÿ“„", - "Image Process ๐Ÿ–ผ๏ธ", - "Test OCR ๐Ÿ”", - "MD Gallery & Process ๐Ÿ“š", - "Build Titan ๐ŸŒฑ", - "Test Image Gen ๐ŸŽจ", + "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 ๐Ÿ–ผ๏ธ" + "Character Gallery ๐Ÿ–ผ๏ธ", ] -tabs = st.tabs(tab_list) -# --- Tab 1: Image -> PDF Layout (from App 1) ------------- +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") - # Camera scan specific to this tab layout_cam = st.camera_input("๐Ÿ“ธ Scan Document for Layout PDF", key="layout_cam") if layout_cam: - central = pytz.timezone("US/Central") # Consider making timezone configurable - now = datetime.now(central) - # Use generate_filename helper + now = datetime.now(pytz.timezone("US/Central")) scan_name = generate_filename(f"layout_scan_{now.strftime('%a').upper()}", "png") - try: - # Save the uploaded file content - with open(scan_name, "wb") as f: - f.write(layout_cam.getvalue()) + 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) + 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}") - # No rerun needed, handled by Streamlit's camera widget update - except Exception as e: - st.error(f"Failed to save scan: {e}") - logger.error(f"Failed to save camera scan {scan_name}: {e}") - - # File uploader specific to this tab - layout_uploads = st.file_uploader( - "๐Ÿ“‚ Upload PNG/JPG Images for Layout PDF", type=["png","jpg","jpeg"], - accept_multiple_files=True, key="layout_uploader" - ) - # Display uploaded images immediately - if layout_uploads: - st.write(f"Uploaded {len(layout_uploads)} images:") - # Keep track of newly uploaded file objects for the DataFrame - st.session_state['layout_new_uploads'] = layout_uploads - + update_gallery(); st.rerun() # 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") - - # --- Build combined list for this tab's purpose --- layout_records = [] - - # From layout-specific snapshots - processed_snapshots = set() # Keep track to avoid duplicates if script reruns + 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, # Store path for snapshots - "width": w, - "height": h, - "aspect_ratio": ar, - "orientation": orient, - "order": idx, # Initial order based on addition - "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)}") - - - # From layout-specific uploads (use the file objects directly) - # Access the newly uploaded files from session state if they exist + 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) # Reset pointer - 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, # Store file object for uploads - "width": w, - "height": h, - "aspect_ratio": ar, - "orientation": orient, - "order": jdx, # Initial order - "type": "Upload" - }) - f_obj.seek(0) # Reset pointer again for potential later use - 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}") + 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.") + if not layout_records: st.info("Scan or upload images using the controls on the left.") else: - # Create DataFrame - layout_df = pd.DataFrame(layout_records) - - # Filter Options (moved here for clarity) - st.markdown("Filter by Orientation:") - dims = st.multiselect( - "Include orientations:", options=["Landscape","Portrait","Square"], - default=["Landscape","Portrait","Square"], key="layout_dims_filter" - ) - if dims: # Apply filter only if options are selected - filtered_df = layout_df[layout_df['orientation'].isin(dims)].copy() # Use copy to avoid SettingWithCopyWarning - else: - filtered_df = layout_df.copy() # No filter applied - - # Ensure 'order' column is integer for editing/sorting - filtered_df['order'] = filtered_df['order'].astype(int) - filtered_df = filtered_df.sort_values('order').reset_index(drop=True) - + 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:") - # Use st.data_editor for reordering - edited_df = st.data_editor( - filtered_df, - column_config={ - "filename": st.column_config.TextColumn("Filename", disabled=True), - "source": None, # Hide source column - "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", # Allow sorting/reordering by drag-and-drop - key="layout_editor" - ) - - # Sort by the edited 'order' column to get the final 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) - - # Extract the sources in the correct order for PDF generation - # Need to handle both file paths (str) and uploaded file objects ordered_sources_for_pdf = ordered_layout_df['source'].tolist() - # --- Generate & Download --- 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.") + if not ordered_sources_for_pdf: st.warning("No images selected or available after filtering.") else: - with st.spinner("Generating PDF... This might take a while for many images."): - pdf_bytes = make_image_sized_pdf(ordered_sources_for_pdf) - + with st.spinner("Generating PDF..."): pdf_bytes = make_image_sized_pdf(ordered_sources_for_pdf) if pdf_bytes: - # Create filename for the PDF - central = pytz.timezone("US/Central") # Use same timezone - now = datetime.now(central) - prefix = now.strftime("%Y%m%d-%H%M%p") - # Create a basename from first few image names - stems = [] - for src in ordered_sources_for_pdf[:4]: # Limit to first 4 - if isinstance(src, str): stems.append(clean_stem(src)) - else: stems.append(clean_stem(getattr(src, 'name', 'upload'))) - basename = " - ".join(stems) - if not basename: basename = "Layout" # Fallback name - pdf_fname = f"{prefix}_{basename}.pdf" - pdf_fname = re.sub(r'[^\w\- \.]', '_', pdf_fname) # Sanitize filename - + 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" - ) - - # Add PDF Preview + 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)) # Standard resolution preview - 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.") + 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 ImportError: - st.info("Install PyMuPDF (`pip install pymupdf`) to enable PDF previews.") - 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.") - + 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.") -# --- Remaining Tabs (from App 2, adapted) ---------------- -# --- Tab: Camera Snap --- +# --- 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: - # Use generate_filename helper - filename = generate_filename("cam0_snap") - # Remove previous file for this camera if it exists - 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 # Ignore error if file is already gone + 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() # Update sidebar gallery - st.rerun() # Rerun to reflect change immediately in 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}") - + 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(); st.rerun() + 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 + 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() # Update sidebar gallery - st.rerun() - 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: Download PDFs --- + 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(); st.rerun() + 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", # Example paper 1 - "https://arxiv.org/pdf/1706.03762", # Attention is All You Need - "https://arxiv.org/pdf/2402.17764", # Example paper 2 - # Add more diverse examples if needed - "https://www.un.org/esa/sustdev/publications/publications.html" # Example non-PDF page (will fail download) - "https://www.clickdimensions.com/links/ACCERL/" # Example direct PDF link - ] + 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" - ) - + 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.") + 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() # Get current list once - + 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 # Still track it - # Ensure it's selectable in the gallery if it exists - if os.path.exists(output_path): - st.session_state['asset_checkboxes'][output_path] = st.session_state['asset_checkboxes'].get(output_path, False) + 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 # Default to unselected - download_count += 1 - existing_pdfs.append(output_path) # Add to current list - else: - st.error(f"Failed to download: {url}") - + 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() # Update sidebar only if new files were added - st.rerun() + if download_count > 0: update_gallery(); st.rerun() 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] - + 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']) # Only get PDFs - if st.session_state['asset_checkboxes'].get(path, False) - ] - - if not selected_pdfs: - st.warning("No PDFs selected in the sidebar gallery! Tick the 'Select' box for PDFs you want to snapshot.") + 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 + 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 - - # Run the async snapshot function - # Need to run asyncio event loop properly in Streamlit + 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) - # Display the generated snapshots - st.write(f"Snapshots for {os.path.basename(pdf_path)}:") - cols = st.columns(3) + 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, unselected - - if total_snapshots_generated > 0: - st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs.") - update_gallery() # Refresh sidebar - st.rerun() - else: - st.warning("No snapshots were generated. Check logs or PDF files.") - + 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(); st.rerun() + else: st.warning("No snapshots were generated. Check logs or PDF files.") -# --- Tab: PDF Process --- +# --- Tab 4: Build Titan (Local Models) --- with tabs[3]: - st.header("PDF Process with GPT ๐Ÿ“„") - st.markdown("Upload PDFs, view pages, and extract text using GPT vision models.") + st.header("Build Titan (Local Models) ๐ŸŒฑ") + st.markdown("Download and save models from Hugging Face Hub for local use.") - if not _openai_available: - st.error("OpenAI features are disabled. Cannot process PDFs with GPT.") + if not _transformers_available: + st.error("Transformers library not available. Cannot download or load local models.") else: - gpt_models = ["gpt-4o", "gpt-4o-mini"] # Add more if needed - selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="pdf_process_gpt_model") - detail_level = st.selectbox("Image Detail Level for GPT", ["auto", "low", "high"], key="pdf_process_detail_level", help="Affects how GPT 'sees' the image. 'high' costs more.") - - uploaded_pdfs_process = st.file_uploader("Upload PDF files to process", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader") - - if uploaded_pdfs_process: - process_button = st.button("Process Uploaded PDFs with GPT", key="process_uploaded_pdfs_gpt") - - if process_button: - combined_text_output = f"# GPT ({selected_gpt_model}) PDF Processing Results\n\n" - total_pages_processed = 0 - output_placeholder = st.container() # Container for dynamic updates - - for pdf_file in uploaded_pdfs_process: - output_placeholder.markdown(f"--- \n### Processing: {pdf_file.name}") - pdf_bytes = pdf_file.read() - temp_pdf_path = f"temp_process_{pdf_file.name}" - - # Save temporary file - with open(temp_pdf_path, "wb") as f: f.write(pdf_bytes) - - try: - doc = fitz.open(temp_pdf_path) - num_pages = len(doc) - output_placeholder.info(f"Found {num_pages} pages. Processing with {selected_gpt_model}...") - - doc_text = f"## File: {pdf_file.name}\n\n" - - for i, page in enumerate(doc): - page_start_time = time.time() - page_placeholder = output_placeholder.empty() - page_placeholder.info(f"Processing Page {i + 1}/{num_pages}...") + build_model_type = st.selectbox( + "Select Model Type", + ["Causal LM", "Vision/Multimodal", "OCR", "Diffusion"], # Added more types + key="build_type_local" + ) - # Generate image from page - pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) # Standard high-res for OCR - img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) + 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." + ) - # Display image being processed - # cols = output_placeholder.columns(2) - # cols[0].image(img, caption=f"Page {i+1}", use_container_width=True) + # 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).") - # Process with GPT - prompt_pdf = "Extract all text content visible on this page. Maintain formatting like paragraphs and lists if possible." - gpt_text = process_image_with_prompt(img, prompt_pdf, model=selected_gpt_model, detail=detail_level) + 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") - doc_text += f"### Page {i + 1}\n\n{gpt_text}\n\n---\n\n" - total_pages_processed += 1 - elapsed_page = int(time.time() - page_start_time) - page_placeholder.success(f"Page {i + 1}/{num_pages} processed in {elapsed_page}s.") - # cols[1].text_area(f"GPT Output (Page {i+1})", gpt_text, height=200, key=f"pdf_gpt_out_{pdf_file.name}_{i}") + 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) - combined_text_output += doc_text - doc.close() + 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 - except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: - output_placeholder.error(f"Error opening PDF {pdf_file.name}: {pdf_err}. Skipping.") - logger.warning(f"Error opening PDF {pdf_file.name}: {pdf_err}") - except Exception as e: - output_placeholder.error(f"Error processing {pdf_file.name}: {str(e)}") - logger.error(f"Error processing PDF {pdf_file.name}: {e}") - finally: - # Clean up temporary file - if os.path.exists(temp_pdf_path): - try: os.remove(temp_pdf_path) - except OSError: pass - - if total_pages_processed > 0: - st.markdown("--- \n### Combined Processing Results") - st.markdown(f"Processed a total of {total_pages_processed} pages.") - st.text_area("Full GPT Output", combined_text_output, height=400, key="combined_pdf_gpt_output") + 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}") - # Save combined output to a file - output_filename = generate_filename("gpt_processed_pdfs", "md") + 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() + # Update sidebar selector + st.rerun() + + 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: - with open(output_filename, "w", encoding="utf-8") as f: - f.write(combined_text_output) - st.success(f"Combined output saved to {output_filename}") - st.markdown(get_download_link(output_filename, "text/markdown", "Download Combined MD"), unsafe_allow_html=True) - # Add to gallery automatically - st.session_state['asset_checkboxes'][output_filename] = False - update_gallery() - except IOError as e: - st.error(f"Failed to save combined output file: {e}") - logger.error(f"Failed to save {output_filename}: {e}") - else: - st.warning("No pages were processed.") + # 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}") + st.rerun() + 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: Image Process --- +# --- Tab 5: PDF Process (HF) --- with tabs[4]: - st.header("Image Process with GPT ๐Ÿ–ผ๏ธ") - st.markdown("Upload images and process them using custom prompts with GPT vision models.") + 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 not _openai_available: - st.error("OpenAI features are disabled. Cannot process images with GPT.") + 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: - gpt_models_img = ["gpt-4o", "gpt-4o-mini"] - selected_gpt_model_img = st.selectbox("Select GPT Model", gpt_models_img, key="img_process_gpt_model") - detail_level_img = st.selectbox("Image Detail Level", ["auto", "low", "high"], key="img_process_detail_level") - prompt_img_process = st.text_area( - "Enter prompt for image processing", - "Describe this image in detail. What is happening? What objects are present?", - key="img_process_prompt_area" - ) + 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_images_process = st.file_uploader( - "Upload image files to process", type=["png", "jpg", "jpeg"], - accept_multiple_files=True, key="image_process_uploader" - ) + 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_images_process: - process_img_button = st.button("Process Uploaded Images with GPT", key="process_uploaded_images_gpt") + 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 process_img_button: - combined_img_text_output = f"# GPT ({selected_gpt_model_img}) Image Processing Results\n\n**Prompt:** {prompt_img_process}\n\n---\n\n" - images_processed_count = 0 - output_img_placeholder = st.container() + 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 img_file in uploaded_images_process: - output_img_placeholder.markdown(f"### Processing: {img_file.name}") - img_placeholder = output_img_placeholder.empty() + 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: - img = Image.open(img_file) - cols_img = output_img_placeholder.columns(2) - cols_img[0].image(img, caption=f"Input: {img_file.name}", use_container_width=True) + 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 - # Process with GPT - gpt_img_text = process_image_with_prompt(img, prompt_img_process, model=selected_gpt_model_img, detail=detail_level_img) + output_placeholder_hf.info(f"Processing {len(pages_to_process)} of {num_pages} pages...") + doc_text = f"## File: {pdf_file.name}\n\n" - cols_img[1].text_area(f"GPT Output", gpt_img_text, height=300, key=f"img_gpt_out_{img_file.name}") - combined_img_text_output += f"## Image: {img_file.name}\n\n{gpt_img_text}\n\n---\n\n" - images_processed_count += 1 - output_img_placeholder.success(f"Processed {img_file.name}.") + 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) - except UnidentifiedImageError: - output_img_placeholder.error(f"Cannot identify image file: {img_file.name}. Skipping.") - logger.warning(f"Cannot identify image file {img_file.name}") - except Exception as e: - output_img_placeholder.error(f"Error processing image {img_file.name}: {str(e)}") - logger.error(f"Error processing image {img_file.name}: {e}") + # 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() - if images_processed_count > 0: - st.markdown("--- \n### Combined Image Processing Results") - st.markdown(f"Processed a total of {images_processed_count} images.") - st.text_area("Full GPT Output (Images)", combined_img_text_output, height=400, key="combined_img_gpt_output") + 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)}") - # Save combined output - output_filename_img = generate_filename("gpt_processed_images", "md") + 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_img, "w", encoding="utf-8") as f: - f.write(combined_img_text_output) - st.success(f"Combined image processing output saved to {output_filename_img}") - st.markdown(get_download_link(output_filename_img, "text/markdown", "Download Combined MD"), unsafe_allow_html=True) - st.session_state['asset_checkboxes'][output_filename_img] = False - update_gallery() - except IOError as e: - st.error(f"Failed to save combined image output file: {e}") - logger.error(f"Failed to save {output_filename_img}: {e}") - else: - st.warning("No images were processed.") + 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: Test OCR --- -with tabs[5]: - st.header("Test OCR with GPT-4o ๐Ÿ”") - st.markdown("Select an image or PDF from the gallery and run GPT-4o OCR.") - if not _openai_available: - st.error("OpenAI features are disabled. Cannot perform OCR.") +# --- 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: - gallery_files_ocr = get_gallery_files(['png', 'jpg', 'jpeg', 'pdf']) + 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}") - if not gallery_files_ocr: - st.warning("No images or PDFs in the gallery. Use Camera Snap or Download PDFs first.") - else: - selected_file_ocr = st.selectbox( - "Select Image or PDF from Gallery for OCR", - options=[""] + gallery_files_ocr, # Add empty option - format_func=lambda x: os.path.basename(x) if x else "Select a file...", - key="ocr_select_file" - ) + combined_img_text_hf += f"## Image: {img_file.name}\n\n{hf_img_text}\n\n---\n\n" + images_processed_hf += 1 - if selected_file_ocr: - st.write(f"Selected: {os.path.basename(selected_file_ocr)}") - file_ext_ocr = os.path.splitext(selected_file_ocr)[1].lower() - image_to_ocr = None - page_info = "" + 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)}") - try: - if file_ext_ocr in ['.png', '.jpg', '.jpeg']: - image_to_ocr = Image.open(selected_file_ocr) - elif file_ext_ocr == '.pdf': - doc = fitz.open(selected_file_ocr) - if len(doc) > 0: - # Use first page for single OCR test - pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) # High-res for OCR - image_to_ocr = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) - page_info = " (Page 1)" - else: - st.warning("Selected PDF is empty.") - doc.close() + 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}") - if image_to_ocr: - st.image(image_to_ocr, caption=f"Image for OCR{page_info}", use_container_width=True) - - if st.button("Run GPT-4o OCR on this Image ๐Ÿš€", key="ocr_run_button"): - output_ocr_file = generate_filename(f"ocr_{os.path.splitext(os.path.basename(selected_file_ocr))[0]}", "txt") - st.session_state['processing']['ocr'] = True # Indicate processing - - # Run async OCR function - ocr_result = asyncio.run(process_gpt4o_ocr(image_to_ocr, output_ocr_file)) - - st.session_state['processing']['ocr'] = False # Clear processing flag - - if ocr_result and not ocr_result.startswith("Error"): - entry = f"OCR Test: {selected_file_ocr}{page_info} -> {output_ocr_file}" - st.session_state['history'].append(entry) - st.text_area("OCR Result", ocr_result, height=300, key="ocr_result_display") - if len(ocr_result) > 10: # Basic check if result seems valid - st.success(f"OCR output saved to {output_ocr_file}") - st.markdown(get_download_link(output_ocr_file, "text/plain", "Download OCR Text"), unsafe_allow_html=True) - # Add txt file to gallery - st.session_state['asset_checkboxes'][output_ocr_file] = False - update_gallery() - else: - st.warning("OCR output seems short or empty; file may not contain useful text.") - if os.path.exists(output_ocr_file): os.remove(output_ocr_file) # Clean up empty file - else: - st.error(f"OCR failed. {ocr_result}") - if os.path.exists(output_ocr_file): os.remove(output_ocr_file) # Clean up failed file - - # Option for multi-page PDF OCR - if file_ext_ocr == '.pdf': - if st.button("Run OCR on All Pages of PDF ๐Ÿ“„๐Ÿš€", key="ocr_all_pages_button"): - st.info("Starting full PDF OCR... This may take time.") - try: - doc = fitz.open(selected_file_ocr) - num_pages_pdf = len(doc) - if num_pages_pdf == 0: - st.warning("PDF is empty.") - else: - full_text_ocr = f"# Full OCR Results for {os.path.basename(selected_file_ocr)}\n\n" - total_pages_ocr_processed = 0 - ocr_output_placeholder = st.container() - - for i in range(num_pages_pdf): - page_ocr_start_time = time.time() - page_ocr_placeholder = ocr_output_placeholder.empty() - page_ocr_placeholder.info(f"OCR - Processing Page {i + 1}/{num_pages_pdf}...") - - pix_ocr = doc[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) - image_page_ocr = Image.frombytes("RGB", [pix_ocr.width, pix_ocr.height], pix_ocr.samples) - output_page_ocr_file = generate_filename(f"ocr_{os.path.splitext(os.path.basename(selected_file_ocr))[0]}_p{i+1}", "txt") - - page_ocr_result = asyncio.run(process_gpt4o_ocr(image_page_ocr, output_page_ocr_file)) - - if page_ocr_result and not page_ocr_result.startswith("Error"): - full_text_ocr += f"## Page {i + 1}\n\n{page_ocr_result}\n\n---\n\n" - entry_page = f"OCR Multi: {selected_file_ocr} Page {i + 1} -> {output_page_ocr_file}" - st.session_state['history'].append(entry_page) - # Don't add individual page txt files to gallery to avoid clutter - if os.path.exists(output_page_ocr_file): os.remove(output_page_ocr_file) - total_pages_ocr_processed += 1 - elapsed_ocr_page = int(time.time() - page_ocr_start_time) - page_ocr_placeholder.success(f"OCR - Page {i + 1}/{num_pages_pdf} done ({elapsed_ocr_page}s).") - else: - page_ocr_placeholder.error(f"OCR failed for Page {i+1}. Skipping.") - full_text_ocr += f"## Page {i + 1}\n\n[OCR FAILED]\n\n---\n\n" - if os.path.exists(output_page_ocr_file): os.remove(output_page_ocr_file) - - doc.close() - - if total_pages_ocr_processed > 0: - md_output_file_ocr = generate_filename(f"full_ocr_{os.path.splitext(os.path.basename(selected_file_ocr))[0]}", "md") - try: - with open(md_output_file_ocr, "w", encoding='utf-8') as f: - f.write(full_text_ocr) - st.success(f"Full PDF OCR complete. Combined output saved to {md_output_file_ocr}") - st.markdown(get_download_link(md_output_file_ocr, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True) - st.session_state['asset_checkboxes'][md_output_file_ocr] = False - update_gallery() - except IOError as e: - st.error(f"Failed to save combined OCR file: {e}") - else: - st.warning("No pages were successfully OCR'd from the PDF.") - - except Exception as e: - st.error(f"Error during full PDF OCR: {e}") - logger.error(f"Full PDF OCR failed for {selected_file_ocr}: {e}") - - except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: - st.error(f"Cannot open PDF {os.path.basename(selected_file_ocr)}: {pdf_err}") - except UnidentifiedImageError: - st.error(f"Cannot identify image file: {os.path.basename(selected_file_ocr)}") - except FileNotFoundError: - st.error(f"File not found: {os.path.basename(selected_file_ocr)}. Refresh the gallery.") - except Exception as e: - st.error(f"An error occurred: {e}") - logger.error(f"Error in OCR tab for {selected_file_ocr}: {e}") -# --- Tab: MD Gallery & Process --- +# --- Tab 7: Text Process (HF) --- with tabs[6]: - st.header("MD & Text File Gallery / GPT Processing ๐Ÿ“š") - st.markdown("View, process, and combine Markdown (.md) and Text (.txt) files from the gallery using GPT.") - - if not _openai_available: - st.error("OpenAI features are disabled. Cannot process text files with GPT.") + 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: - gpt_models_md = ["gpt-4o", "gpt-4o-mini"] - selected_gpt_model_md = st.selectbox("Select GPT Model for Text Processing", gpt_models_md, key="md_process_gpt_model") + 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.") - md_txt_files = get_gallery_files(['md', 'txt']) - - if not md_txt_files: - st.warning("No Markdown (.md) or Text (.txt) files found in the gallery.") - else: - st.subheader("Individual File Processing") - selected_file_md = st.selectbox( - "Select MD/TXT File to Process", - options=[""] + md_txt_files, - format_func=lambda x: os.path.basename(x) if x else "Select a file...", - key="md_select_individual" - ) - - if selected_file_md: - st.write(f"Selected: {os.path.basename(selected_file_md)}") - try: - with open(selected_file_md, "r", encoding="utf-8", errors='ignore') as f: - content_md = f.read() - st.text_area("File Content Preview", content_md[:1000] + ("..." if len(content_md) > 1000 else ""), height=200, key="md_content_preview") - - prompt_md_individual = st.text_area( - "Enter Prompt for this File", - "Summarize the key points of this text into a bulleted list.", - key="md_individual_prompt" - ) + 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 st.button(f"Process {os.path.basename(selected_file_md)} with GPT", key=f"process_md_ind_{selected_file_md}"): - with st.spinner("Processing text with GPT..."): - result_text_md = process_text_with_prompt(content_md, prompt_md_individual, model=selected_gpt_model_md) - - st.markdown("### GPT Processing Result") - st.markdown(result_text_md) # Display result as Markdown - - # Save the result - output_filename_md = generate_filename(f"gpt_processed_{os.path.splitext(os.path.basename(selected_file_md))[0]}", "md") - try: - with open(output_filename_md, "w", encoding="utf-8") as f: - f.write(result_text_md) - st.success(f"Processing result saved to {output_filename_md}") - st.markdown(get_download_link(output_filename_md, "text/markdown", "Download Processed MD"), unsafe_allow_html=True) - st.session_state['asset_checkboxes'][output_filename_md] = False - update_gallery() - except IOError as e: - st.error(f"Failed to save processed MD file: {e}") - - except FileNotFoundError: - st.error("Selected file not found. It might have been deleted.") - except Exception as e: - st.error(f"Error reading or processing file: {e}") - - st.markdown("---") - st.subheader("Combine and Process Multiple Files") - st.write("Select MD/TXT files from the gallery to combine:") - - selected_md_combine = {} - cols_md = st.columns(3) - for idx, md_file in enumerate(md_txt_files): - with cols_md[idx % 3]: - selected_md_combine[md_file] = st.checkbox( - f"{os.path.basename(md_file)}", - key=f"checkbox_md_combine_{md_file}" - ) - - prompt_md_combine = st.text_area( - "Enter Prompt for Combined Content", - "Synthesize the following texts into a cohesive summary. Identify the main themes and provide supporting details from the different sources.", - key="md_combine_prompt" - ) + 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}") - if st.button("Process Selected MD/TXT Files with GPT", key="process_combine_md"): - files_to_combine = [f for f, selected in selected_md_combine.items() if selected] + except FileNotFoundError: st.error("Selected file not found.") + except Exception as e: st.error(f"Error reading file: {e}") - if not files_to_combine: - st.warning("No files selected for combination.") - else: - st.info(f"Combining {len(files_to_combine)} files...") - combined_content = "" - for md_file in files_to_combine: - try: - with open(md_file, "r", encoding="utf-8", errors='ignore') as f: - combined_content += f"\n\n## --- Source: {os.path.basename(md_file)} ---\n\n" + f.read() - except Exception as e: - st.error(f"Error reading {md_file}: {str(e)}. Skipping.") - logger.warning(f"Error reading {md_file} for combination: {e}") - - if combined_content: - st.text_area("Preview Combined Content (First 2000 chars)", combined_content[:2000]+"...", height=200) - with st.spinner("Processing combined text with GPT..."): - result_text_combine = process_text_with_prompt(combined_content, prompt_md_combine, model=selected_gpt_model_md) - - st.markdown("### Combined Processing Result") - st.markdown(result_text_combine) - - # Save the combined result - output_filename_combine = generate_filename("gpt_combined_md_txt", "md") - try: - with open(output_filename_combine, "w", encoding="utf-8") as f: - f.write(f"# Combined Processing Result\n\n**Prompt:** {prompt_md_combine}\n\n**Sources:** {', '.join([os.path.basename(f) for f in files_to_combine])}\n\n---\n\n{result_text_combine}") - st.success(f"Combined processing result saved to {output_filename_combine}") - st.markdown(get_download_link(output_filename_combine, "text/markdown", "Download Combined MD"), unsafe_allow_html=True) - st.session_state['asset_checkboxes'][output_filename_combine] = False - update_gallery() - except IOError as e: - st.error(f"Failed to save combined processed file: {e}") - else: - st.error("Failed to read content from selected files.") -# --- Tab: Build Titan --- +# --- Tab 8: Test OCR (HF) --- with tabs[7]: - st.header("Build Titan Model ๐ŸŒฑ") - st.markdown("Download and save base models for Causal LM or Diffusion tasks.") - - if not _ai_libs_available: - st.error("AI/ML libraries (torch, transformers, diffusers) are required for this feature.") + 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: - build_model_type = st.selectbox("Model Type to Build", ["Causal LM", "Diffusion"], key="build_type_select") - - if build_model_type == "Causal LM": - default_causal = "HuggingFaceTB/SmolLM-135M" #"Qwen/Qwen1.5-0.5B-Chat" is larger - causal_models = [default_causal, "gpt2", "distilgpt2"] # Add more small options - base_model_select = st.selectbox( - "Select Base Causal LM", causal_models, index=causal_models.index(default_causal), - key="causal_model_select" - ) - else: # Diffusion - default_diffusion = "OFA-Sys/small-stable-diffusion-v0" #"stabilityai/stable-diffusion-2-base" is large - diffusion_models = [default_diffusion, "google/ddpm-cat-256", "google/ddpm-celebahq-256"] # Add more small options - base_model_select = st.selectbox( - "Select Base Diffusion Model", diffusion_models, index=diffusion_models.index(default_diffusion), - key="diffusion_model_select" - ) - - model_name_build = st.text_input("Local Model Name", f"{build_model_type.lower().replace(' ','')}-titan-{os.path.basename(base_model_select)}-{int(time.time()) % 10000}", key="build_model_name") - domain_build = st.text_input("Optional: Target Domain Tag", "general", key="build_domain") - - if st.button(f"Download & Save {build_model_type} Model โฌ‡๏ธ", key="download_build_model"): - if not model_name_build: - st.error("Please provide a local model name.") - else: - if build_model_type == "Causal LM": - config = ModelConfig( - name=model_name_build, base_model=base_model_select, size="small", domain=domain_build # Size is illustrative - ) - builder = ModelBuilder() - else: - config = DiffusionConfig( - name=model_name_build, base_model=base_model_select, size="small", domain=domain_build - ) - builder = DiffusionBuilder() - - try: - builder.load_model(base_model_select, config) - builder.save_model(config.model_path) # Save to ./models/ or ./diffusion_models/ - - st.session_state['builder'] = builder # Store the loaded builder instance - st.session_state['model_loaded'] = True - st.session_state['selected_model_type'] = build_model_type - st.session_state['selected_model'] = config.model_path # Store path to local copy + 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" + ) - st.session_state['history'].append(f"Built {build_model_type} model: {model_name_build} from {base_model_select}") - st.success(f"{build_model_type} model downloaded from {base_model_select} and saved locally to {config.model_path}! ๐ŸŽ‰") - # No automatic rerun, let user proceed - except Exception as e: - st.error(f"Failed to build model: {e}") - logger.error(f"Failed to build model {model_name_build} from {base_model_select}: {e}") + 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 = "" -# --- Tab: Test Image Gen --- + 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 ๐ŸŽจ") - st.markdown("Generate images using a loaded Diffusion model.") + st.header("Test Image Generation (Diffusers) ๐ŸŽจ") + st.markdown("Generate images using Stable Diffusion models loaded locally via the Diffusers library.") - if not _ai_libs_available: - st.error("AI/ML libraries (torch, transformers, diffusers) are required for image generation.") + if not _diffusers_available: + st.error("Diffusers library is required for image generation.") else: - # Check if a diffusion model is loaded in session state or select one - available_diffusion_models = get_model_files("diffusion") - loaded_diffusion_model_path = None - - # Check if the currently loaded builder is diffusion - current_builder = st.session_state.get('builder') - if current_builder and isinstance(current_builder, DiffusionBuilder) and current_builder.pipeline: - loaded_diffusion_model_path = current_builder.config.model_path if current_builder.config else "Loaded Model" - - # Prepare options for selection, prioritizing loaded model - model_options = ["Load Default Small Model"] + available_diffusion_models - current_selection_index = 0 # Default to loading small model - if loaded_diffusion_model_path and loaded_diffusion_model_path != "Loaded Model": - if loaded_diffusion_model_path not in model_options: - model_options.insert(1, loaded_diffusion_model_path) # Add if not already listed - current_selection_index = model_options.index(loaded_diffusion_model_path) - elif loaded_diffusion_model_path == "Loaded Model": - # A model is loaded, but we don't have its path (e.g., loaded directly) - model_options.insert(1, "Currently Loaded Model") - current_selection_index = 1 - - - selected_diffusion_model = st.selectbox( - "Select Diffusion Model for Generation", - options=model_options, - index=current_selection_index, - key="imggen_model_select", - help="Select a locally saved model, or load the default small one." - ) + # 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") - # Button to explicitly load the selected model if it's not the active one - load_needed = False - if selected_diffusion_model == "Load Default Small Model": - load_needed = not (current_builder and isinstance(current_builder, DiffusionBuilder) and current_builder.config and current_builder.config.base_model == "OFA-Sys/small-stable-diffusion-v0") - elif selected_diffusion_model == "Currently Loaded Model": - load_needed = False # Already loaded - else: # A specific path is selected - load_needed = not (current_builder and isinstance(current_builder, DiffusionBuilder) and current_builder.config and current_builder.config.model_path == selected_diffusion_model) - - if load_needed: - if st.button(f"Load '{os.path.basename(selected_diffusion_model)}' Model", key="imggen_load_sel"): + 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: - if selected_diffusion_model == "Load Default Small Model": - model_to_load = "OFA-Sys/small-stable-diffusion-v0" - config = DiffusionConfig(name="default-small", base_model=model_to_load, size="small") - builder = DiffusionBuilder().load_model(model_to_load, config) - st.session_state['builder'] = builder - st.session_state['model_loaded'] = True - st.session_state['selected_model_type'] = "Diffusion" - st.session_state['selected_model'] = config.model_path # This isn't saved, just track base - st.success("Default small diffusion model loaded.") - st.rerun() - else: # Load from local path - config = DiffusionConfig(name=os.path.basename(selected_diffusion_model), base_model="local", size="unknown", model_path=selected_diffusion_model) - builder = DiffusionBuilder().load_model(selected_diffusion_model, config) - st.session_state['builder'] = builder - st.session_state['model_loaded'] = True - st.session_state['selected_model_type'] = "Diffusion" - st.session_state['selected_model'] = config.model_path - st.success(f"Loaded local model: {config.name}") - st.rerun() - except Exception as e: - st.error(f"Failed to load model {selected_diffusion_model}: {e}") - logger.error(f"Failed loading diffusion model {selected_diffusion_model}: {e}") + # 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!") - # Image Generation Prompt - prompt_imggen = st.text_area("Image Generation Prompt", "A futuristic cityscape at sunset, neon lights, flying cars", key="imggen_prompt") + # 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}") - if st.button("Generate Image ๐Ÿš€", key="imggen_run_button"): - # Check again if a model is effectively loaded and ready - current_builder = st.session_state.get('builder') - if not (current_builder and isinstance(current_builder, DiffusionBuilder) and current_builder.pipeline): - st.error("No diffusion model is loaded. Please select and load a model first.") - elif not prompt_imggen: - st.warning("Please enter a prompt.") - else: - output_imggen_file = generate_filename("image_gen", "png") - st.session_state['processing']['gen'] = True - - # Run async generation - generated_image = asyncio.run(process_image_gen(prompt_imggen, output_imggen_file)) - - st.session_state['processing']['gen'] = False - - if generated_image and os.path.exists(output_imggen_file): - entry = f"Image Gen: '{prompt_imggen[:30]}...' -> {output_imggen_file}" - st.session_state['history'].append(entry) - st.image(generated_image, caption=f"Generated: {os.path.basename(output_imggen_file)}", use_container_width=True) - st.success(f"Image saved to {output_imggen_file}") - st.markdown(get_download_link(output_imggen_file, "image/png", "Download Generated Image"), unsafe_allow_html=True) - # Add to gallery - st.session_state['asset_checkboxes'][output_imggen_file] = False - update_gallery() - # Consider st.rerun() if immediate gallery update is critical - else: - st.error("Image generation failed. Check logs.") + 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: Character Editor --- +# --- 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 or Modify Your Character") - - # Load existing characters for potential editing (optional) - load_characters() - existing_char_names = [c['name'] for c in st.session_state.get('characters', [])] - - # Use a unique key for the form to allow reset + 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**") - # Randomize button inside the form - if st.form_submit_button("Randomize Content ๐ŸŽฒ"): - # Increment key to force form reset with new random values - st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1 - st.rerun() # Rerun to get new random defaults in the reset form - - # Get random defaults only once per form rendering cycle unless reset + 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, letters, numbers, underscore, hyphen, space)", - 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" - ) + 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: - # Validation error = False - if not (3 <= len(name_char) <= 25): - st.error("Name must be between 3 and 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"Character name '{name_char}' already exists!") - error = True - if not intro_char or not greeting_char: - st.error("Intro and Greeting cannot be empty.") - error = True - + 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'), # Added timezone - "tags": tag_list - } + 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 successfully!") - # Increment key to reset form for next creation - st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1 - st.rerun() # Rerun to clear form and update gallery tab - + 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: Character Gallery --- +# --- Tab 11: Character Gallery (Keep from previous merge) --- with tabs[10]: + # ... (Code from previous merge for this tab) ... st.header("Character Gallery ๐Ÿ–ผ๏ธ") - - # Load characters every time the tab is viewed - load_characters() - characters_list = st.session_state.get('characters', []) - - if not characters_list: - st.warning("No characters created yet. Use the Character Editor tab!") + 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)})") - st.markdown("View and manage your created characters.") - - # Search/Filter (Optional Enhancement) 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) # Adjust number of columns as needed - chars_to_delete = [] # Store names to delete after iteration - + 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]: - with st.container(border=True): # Add border to each character card - st.markdown(f"**{char['name']}**") - st.caption(f"Gender: {char.get('gender', 'N/A')}") # Use .get for safety - st.markdown("**Intro:**") - st.markdown(f"> {char.get('intro', '')}") # Blockquote style - 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 Button - delete_key_char = f"delete_char_{char['name']}_{idx}" # More unique key - if st.button(f"Delete {char['name']}", key=delete_key_char, type="primary"): - chars_to_delete.append(char['name']) # Mark for deletion - - # Process deletions after iterating + 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] + 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() # Rerun to reflect changes - except IOError as e: - logger.error(f"Failed to save characters.json after deletion: {e}") - st.error("Failed to update character file after deletion.") - + 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 ------------ @@ -1827,17 +1585,16 @@ update_gallery() 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:]]) # Show last 20 logs + 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: - # Display history in reverse chronological order for entry in reversed(st.session_state.get("history", [])): if entry: history_expander.write(f"- {entry}") st.sidebar.markdown("---") -st.sidebar.info("App combines Image Layout PDF generation with AI Vision/SFT tools.") -st.sidebar.caption("Combined App by AI Assistant for User") \ No newline at end of file +st.sidebar.info("Using Hugging Face models for AI tasks.") +st.sidebar.caption("App Modified by AI Assistant") \ No newline at end of file