diff --git "a/app.py" "b/app.py" --- "a/app.py" +++ "b/app.py" @@ -1,3 +1,4 @@ +# --- Combined Imports ------------------------------------ import io import os import re @@ -11,76 +12,65 @@ import zipfile import json import asyncio import aiofiles -import toml + from datetime import datetime from collections import Counter from dataclasses import dataclass, field from io import BytesIO from typing import Optional, List, Dict, Any + import pandas as pd import pytz import streamlit as st -from PIL import Image, ImageDraw +from PIL import Image, ImageDraw # Added ImageDraw from reportlab.pdfgen import canvas from reportlab.lib.utils import ImageReader -from reportlab.lib.pagesizes import letter -from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak -from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle -from reportlab.lib.enums import TA_JUSTIFY -import fitz -import requests +from reportlab.lib.pagesizes import letter # Default page size +import fitz # PyMuPDF + +# --- Hugging Face Imports --- +from huggingface_hub import InferenceClient, HfApi, list_models +from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Import specific exceptions + + + +# --- App Configuration ----------------------------------- +st.set_page_config( + page_title="Vision & Layout Titans (HF) πŸš€πŸ–ΌοΈ", + page_icon="πŸ€–", + layout="wide", + initial_sidebar_state="expanded", + menu_items={ + 'Get Help': 'https://huggingface.co/docs', + 'Report a Bug': None, # Replace with your bug report link if desired + 'About': "Combined App: Image->PDF Layout + Hugging Face Powered AI Tools 🌌" + } +) + + +# Conditional imports for optional/heavy libraries try: import torch - from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, pipeline + from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline + # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.) _transformers_available = True except ImportError: _transformers_available = False st.sidebar.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.") + try: from diffusers import StableDiffusionPipeline _diffusers_available = True except ImportError: _diffusers_available = False + # Don't show warning if transformers also missing, handled above if _transformers_available: st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.") -try: - from openai import OpenAI - _openai_available = True -except ImportError: - _openai_available = False - st.sidebar.warning("OpenAI library not found. OpenAI model features disabled.") -from huggingface_hub import InferenceClient, HfApi, list_models -from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError -# --- App Configuration --- -st.set_page_config( - page_title="Vision & Layout Titans πŸš€πŸ–ΌοΈ", - page_icon="πŸ€–", - layout="wide", - initial_sidebar_state="expanded", - menu_items={ - 'Get Help': 'https://huggingface.co/docs', - 'Report a Bug': None, - 'About': "Combined App: Image/MD->PDF Layout + AI-Powered Tools 🌌" - } -) -# --- Secrets Management --- -try: - secrets = toml.load(".streamlit/secrets.toml") if os.path.exists(".streamlit/secrets.toml") else {} - HF_TOKEN = secrets.get("HF_TOKEN", os.getenv("HF_TOKEN", "")) - OPENAI_API_KEY = secrets.get("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY", "")) -except Exception as e: - st.error(f"Error loading secrets: {e}") - HF_TOKEN = os.getenv("HF_TOKEN", "") - OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") - -if not HF_TOKEN: - st.sidebar.warning("Hugging Face token not found in secrets or environment. Some features may be limited.") -if not OPENAI_API_KEY and _openai_available: - st.sidebar.warning("OpenAI API key not found in secrets or environment. OpenAI features disabled.") - -# --- Logging Setup --- +import requests # Keep requests import + +# --- Logging Setup --------------------------------------- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) log_records = [] @@ -89,91 +79,95 @@ class LogCaptureHandler(logging.Handler): log_records.append(record) logger.addHandler(LogCaptureHandler()) -# --- Model Initialization --- +# --- 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", + "meta-llama/Meta-Llama-3.1-8B-Instruct", # Updated Llama model "mistralai/Mistral-7B-Instruct-v0.3", - "google/gemma-2-9b-it", - "Qwen/Qwen2-7B-Instruct", + "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", + "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", - "llava-hf/llava-1.5-7b-hf", - "google/vit-base-patch16-224" + "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", - "runwayml/stable-diffusion-v1-5", - "OFA-Sys/small-stable-diffusion-v0" + "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 ] -OPENAI_MODELS_LIST = [ - "gpt-4o", - "gpt-4-turbo", - "gpt-3.5-turbo", - "text-davinci-003" -] -st.session_state.setdefault('local_models', {}) -st.session_state.setdefault('hf_inference_client', None) -st.session_state.setdefault('openai_client', None) -if _openai_available and OPENAI_API_KEY: - try: - st.session_state['openai_client'] = OpenAI(api_key=OPENAI_API_KEY) - logger.info("OpenAI client initialized successfully.") - except Exception as e: - st.error(f"Failed to initialize OpenAI client: {e}") - logger.error(f"OpenAI client initialization failed: {e}") - st.session_state['openai_client'] = None -# --- Session State Initialization --- + +# --- Session State Initialization (Combined & Updated) --- +# Layout PDF specific st.session_state.setdefault('layout_snapshots', []) st.session_state.setdefault('layout_new_uploads', []) + +# General App State st.session_state.setdefault('history', []) st.session_state.setdefault('processing', {}) -st.session_state.setdefault('asset_checkboxes', {'image': {}, 'md': {}, 'pdf': {}}) +st.session_state.setdefault('asset_checkboxes', {}) st.session_state.setdefault('downloaded_pdfs', {}) st.session_state.setdefault('unique_counter', 0) st.session_state.setdefault('cam0_file', None) st.session_state.setdefault('cam1_file', None) st.session_state.setdefault('characters', []) -st.session_state.setdefault('char_form_reset_key', 0) +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]) -st.session_state.setdefault('hf_custom_api_model', "") -st.session_state.setdefault('openai_selected_model', OPENAI_MODELS_LIST[0] if _openai_available else "") -st.session_state.setdefault('selected_local_model_path', None) +st.session_state.setdefault('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'] = {'image': st.sidebar.empty(), 'md': st.sidebar.empty(), 'pdf': st.sidebar.empty()} + st.session_state['asset_gallery_container'] = st.sidebar.empty() -# --- Dataclasses --- +# --- Dataclasses (Refined for Local Models) ------------- @dataclass class LocalModelConfig: - name: str - hf_id: str - model_type: str - size_category: str = "unknown" + name: str # User-defined local name + hf_id: str # Hugging Face model ID used for download + model_type: str # 'causal', 'vision', 'diffusion', 'ocr', etc. + size_category: str = "unknown" # e.g., 'small', 'medium', 'large' domain: Optional[str] = None - local_path: str = field(init=False) + 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) + 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 @@ -182,7 +176,10 @@ class DiffusionConfig: def model_path(self): return f"diffusion_models/{self.name}" -# --- Helper Functions --- + +# --- Helper Functions (Combined and refined) ------------- +# (Keep generate_filename, pdf_url_to_filename, get_download_link, zip_directory) +# ... (previous helper functions like generate_filename, pdf_url_to_filename etc. are assumed here) ... def generate_filename(sequence, ext="png"): timestamp = time.strftime('%Y%m%d_%H%M%S') safe_sequence = re.sub(r'[^\w\-]+', '_', str(sequence)) @@ -191,14 +188,12 @@ def generate_filename(sequence, ext="png"): def pdf_url_to_filename(url): name = re.sub(r'^https?://', '', url) name = re.sub(r'[<>:"/\\|?*]', '_', name) - return name[:100] + ".pdf" + return name[:100] + ".pdf" # Limit length def get_download_link(file_path, mime_type="application/octet-stream", label="Download"): - if not os.path.exists(file_path): - return f"{label} (File not found)" + 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: @@ -213,6 +208,7 @@ def zip_directory(directory_path, zip_path): zipf.write(file_path, os.path.relpath(file_path, os.path.dirname(directory_path))) def get_local_model_paths(model_type="causal"): + """Gets paths of locally saved models of a specific type.""" pattern = f"{model_type}_models/*" dirs = [d for d in glob.glob(pattern) if os.path.isdir(d)] return dirs @@ -222,16 +218,10 @@ def get_gallery_files(file_types=("png", "pdf", "jpg", "jpeg", "md", "txt")): for ext in file_types: all_files.update(glob.glob(f"*.{ext.lower()}")) all_files.update(glob.glob(f"*.{ext.upper()}")) - return sorted([f for f in all_files if os.path.basename(f).lower() != 'readme.md']) + return sorted(list(all_files)) -def get_typed_gallery_files(file_type): - if file_type == 'image': - return get_gallery_files(('png', 'jpg', 'jpeg')) - elif file_type == 'md': - return get_gallery_files(('md',)) - elif file_type == 'pdf': - return get_gallery_files(('pdf',)) - return [] +def get_pdf_files(): + return sorted(glob.glob("*.pdf") + glob.glob("*.PDF")) def download_pdf(url, output_path): try: @@ -239,27 +229,27 @@ def download_pdf(url, output_path): 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}") - if os.path.exists(output_path): - try: - os.remove(output_path) - except: + 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: + if os.path.exists(output_path): + try: + os.remove(output_path) + except: pass return False +# (Keep process_pdf_snapshot - it doesn't use AI) async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): start_time = time.time() status_placeholder = st.empty() @@ -268,7 +258,11 @@ async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): try: doc = fitz.open(pdf_path) matrix = fitz.Matrix(resolution_factor, resolution_factor) - num_pages_to_process = min(1, len(doc)) if mode == "single" else min(2, len(doc)) if mode == "twopage" else len(doc) + num_pages_to_process = 0 + if mode == "single": num_pages_to_process = min(1, len(doc)) + elif mode == "twopage": num_pages_to_process = min(2, len(doc)) + elif mode == "allpages": num_pages_to_process = len(doc) + for i in range(num_pages_to_process): page_start_time = time.time() page = doc[i] @@ -280,6 +274,7 @@ async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): elapsed_page = int(time.time() - page_start_time) status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... Page {i+1}/{num_pages_to_process} done ({elapsed_page}s)") await asyncio.sleep(0.01) + doc.close() elapsed = int(time.time() - start_time) status_placeholder.success(f"PDF Snapshot ({mode}, {len(output_files)} files) completed in {elapsed}s!") @@ -288,288 +283,300 @@ async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0): logger.error(f"Failed to process PDF snapshot for {pdf_path}: {e}") status_placeholder.error(f"Failed to process PDF {os.path.basename(pdf_path)}: {e}") for f in output_files: - if os.path.exists(f): - os.remove(f) + if os.path.exists(f): os.remove(f) return [] + +# --- HF Inference Client Management --- def get_hf_client() -> Optional[InferenceClient]: + """Gets or initializes the Hugging Face Inference Client based on session state.""" provider = st.session_state.hf_provider custom_key = st.session_state.hf_custom_key.strip() token_to_use = custom_key if custom_key else HF_TOKEN + if not token_to_use and provider != "hf-inference": - st.error(f"Provider '{provider}' requires a Hugging Face API token.") + 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 + 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: - 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 + # 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}.") + 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: {e}") + st.error(f"Failed to initialize Hugging Face client for provider {provider}: {e}") logger.error(f"InferenceClient initialization failed: {e}") st.session_state.hf_inference_client = None + return st.session_state.hf_inference_client -def process_text_hf(text: str, prompt: str, use_api: bool, model_id: str = None) -> str: +# --- HF/Local Model Processing Functions (Replaced OpenAI ones) --- + +def process_text_hf(text: str, prompt: str, use_api: bool) -> str: + """Processes text using either HF Inference API or a loaded local model.""" status_placeholder = st.empty() start_time = time.time() result_text = "" + + # --- Prepare Parameters --- params = { - "max_new_tokens": st.session_state.gen_max_tokens, + "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, + "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 - system_prompt = "You are a helpful assistant. Process the following text based on the user's request." + 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." - model_id = model_id or st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model + 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'], + max_tokens=params['max_new_tokens'], # chat_completion uses max_tokens temperature=params['temperature'], top_p=params['top_p'], + # Add other params if supported by client.chat_completion ) result_text = response.choices[0].message.content or "" logger.info(f"HF API text processing successful for model {model_id}.") + except Exception as e: logger.error(f"HF API text processing failed for model {model_id}: {e}") result_text = f"Error during Hugging Face API inference: {str(e)}" + else: + # --- Use Loaded Local Model --- status_placeholder.info("Processing text using local model...") if not _transformers_available: - return "Error: Transformers library not available." + 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." + 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." + 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)}." + return f"Error: Model or tokenizer not found for {os.path.basename(model_path)}." + try: + # Prepare input for local transformers model + # Handle chat template if available, otherwise basic concatenation try: prompt_for_model = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - except Exception: - 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) + 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), - "do_sample": True if params['temperature'] > 0.1 else False, - "pad_token_id": tokenizer.eos_token_id + "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 } - with torch.no_grad(): - outputs = model.generate(**inputs, **generate_params) + if 'seed' in params: pass # Seed handling can be complex with transformers, often set globally + + with torch.no_grad(): # Disable gradient calculation for inference + outputs = model.generate(**inputs, **generate_params) + + # Decode the output, skipping special tokens and the prompt + # output_text = tokenizer.decode(outputs[0], skip_special_tokens=True) + # More robust decoding: only decode the newly generated part input_length = inputs['input_ids'].shape[1] generated_ids = outputs[0][input_length:] result_text = tokenizer.decode(generated_ids, skip_special_tokens=True) + logger.info(f"Local text processing successful for model {model_path}.") + except Exception as e: logger.error(f"Local text processing failed for model {model_path}: {e}") result_text = f"Error during local model inference: {str(e)}" - elapsed = int(time.time() - start_time) - status_placeholder.success(f"Text processing completed in {elapsed}s.") - return result_text -def process_text_openai(text: str, prompt: str, model_id: str) -> str: - if not _openai_available or not st.session_state.get('openai_client'): - return "Error: OpenAI client not available or API key missing." - status_placeholder = st.empty() - start_time = time.time() - client = st.session_state['openai_client'] - system_prompt = "You are a helpful assistant. Process the following text based on the user's request." - full_prompt = f"{prompt}\n\n---\n\n{text}" - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": full_prompt} - ] - status_placeholder.info(f"Processing text using OpenAI model: {model_id}...") - try: - response = client.chat.completions.create( - model=model_id, - messages=messages, - max_tokens=st.session_state.gen_max_tokens, - temperature=st.session_state.gen_temperature, - top_p=st.session_state.gen_top_p, - ) - result_text = response.choices[0].message.content or "" - logger.info(f"OpenAI text processing successful for model {model_id}.") - except Exception as e: - logger.error(f"OpenAI text processing failed for model {model_id}: {e}") - result_text = f"Error during OpenAI inference: {str(e)}" + elapsed = int(time.time() - start_time) status_placeholder.success(f"Text processing completed in {elapsed}s.") return result_text -def process_image_hf(image: Image.Image, prompt: str, use_api: bool, model_id: str = None) -> str: + +# --- Image Processing (Placeholder/Basic Implementation) --- +# This needs significant work depending on the chosen vision model type +def process_image_hf(image: Image.Image, prompt: str, use_api: bool) -> str: + """Processes an image using either HF Inference API or a local model.""" status_placeholder = st.empty() start_time = time.time() - result_text = "" + result_text = "[Image processing not fully implemented with HF models yet]" + if use_api: - status_placeholder.info("Processing image using Hugging Face 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." + 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() - model_id = model_id or "Salesforce/blip-image-captioning-large" - status_placeholder.info(f"Using API Image-to-Text Model: {model_id}") + try: - response_list = client.image_to_text(data=img_bytes, model=model_id) + # 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 = response_list[0]['generated_text'] - logger.info(f"HF API image captioning successful for model {model_id}.") + 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}") + 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." + 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." + 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') - if model_type == 'vision': - processor = local_model_data.get('processor') - model = local_model_data.get('model') - if processor and model: - try: - inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device) - generated_ids = model.generate(**inputs, max_new_tokens=st.session_state.gen_max_tokens) - result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() - except Exception as e: - result_text = f"Error during local vision model inference: {e}" - else: - result_text = "Error: Processor or model missing for local vision task." - elif model_type == 'ocr': - processor = local_model_data.get('processor') - model = local_model_data.get('model') - if processor and model: - try: - pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(model.device) - generated_ids = model.generate(pixel_values, max_new_tokens=st.session_state.gen_max_tokens) - result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] - except Exception as e: - result_text = f"Error during local OCR model inference: {e}" - else: - result_text = "Error: Processor or model missing for local OCR task." + + # --- 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." - elapsed = int(time.time() - start_time) - status_placeholder.success(f"Image processing completed in {elapsed}s.") - return result_text + result_text = f"Error: Loaded model '{os.path.basename(model_path)}' is not a recognized vision/OCR type for this function." + -def process_image_openai(image: Image.Image, prompt: str, model_id: str = "gpt-4o") -> str: - if not _openai_available or not st.session_state.get('openai_client'): - return "Error: OpenAI client not available or API key missing." - status_placeholder = st.empty() - start_time = time.time() - client = st.session_state['openai_client'] - buffered = BytesIO() - image.save(buffered, format="PNG") - img_b64 = base64.b64encode(buffered.getvalue()).decode() - status_placeholder.info(f"Processing image using OpenAI model: {model_id}...") - try: - response = client.chat.completions.create( - model=model_id, - messages=[ - {"role": "user", "content": [ - {"type": "text", "text": prompt}, - {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}} - ]} - ], - max_tokens=st.session_state.gen_max_tokens, - temperature=st.session_state.gen_temperature, - ) - result_text = response.choices[0].message.content or "" - logger.info(f"OpenAI image processing successful for model {model_id}.") - except Exception as e: - logger.error(f"OpenAI image processing failed for model {model_id}: {e}") - result_text = f"Error during OpenAI image inference: {str(e)}" elapsed = int(time.time() - start_time) - status_placeholder.success(f"Image processing completed in {elapsed}s.") + status_placeholder.success(f"Image processing attempt completed in {elapsed}s.") return result_text -async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool, model_id: str = None) -> str: +# 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, model_id=model_id or "microsoft/trocr-large-handwritten") + result = process_image_hf(image, ocr_prompt, use_api) + + # Save the result if it looks like text (basic check) if result and not result.startswith("Error") and not result.startswith("["): try: async with aiofiles.open(output_file, "w", encoding='utf-8') as f: await f.write(result) logger.info(f"HF OCR result saved to {output_file}") except IOError as e: - logger.error(f"Failed to save HF OCR output to {output_file}: {e}") - result += f"\n[Error saving file: {e}]" + 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): - try: - os.remove(output_file) - except OSError: - pass - return result + # Remove file if processing failed or was just a placeholder message + try: os.remove(output_file) + except OSError: pass -async def process_openai_ocr(image: Image.Image, output_file: str, model_id: str = "gpt-4o") -> str: - ocr_prompt = "Extract text content from this image." - result = process_image_openai(image, ocr_prompt, model_id) - if result and not result.startswith("Error"): - try: - async with aiofiles.open(output_file, "w", encoding='utf-8') as f: - await f.write(result) - logger.info(f"OpenAI OCR result saved to {output_file}") - except IOError as e: - logger.error(f"Failed to save OpenAI OCR output to {output_file}: {e}") - result += f"\n[Error saving file: {e}]" - elif os.path.exists(output_file): - try: - os.remove(output_file) - except OSError: - pass return result + +# --- Character Functions (Keep from previous) ----------- +# ... (randomize_character_content, save_character, load_characters are assumed here) ... def randomize_character_content(): intro_templates = [ "{char} is a valiant knight...", "{char} is a mischievous thief...", @@ -578,8 +585,7 @@ def randomize_character_content(): greeting_templates = [ "'I am from the knight's guild...'", "'I heard you needed helpβ€”name’s {char}...", "'Oh, hello! I’m {char}, didn’t see you there...'", "'I’m {char}, and I’m here to fight...'", - "'I’m {char}, here to heal...'" - ] + "'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) @@ -589,13 +595,12 @@ def randomize_character_content(): def save_character(character_data): characters = st.session_state.get('characters', []) if any(c['name'] == character_data['name'] for c in characters): - st.error(f"Character name '{character_data['name']}' already exists.") - return False + 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) + 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: @@ -604,336 +609,1085 @@ def save_character(character_data): return False def load_characters(): - 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) - 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") + with open("characters.json", "r", encoding='utf-8') as f: characters = json.load(f) + if isinstance(characters, list): st.session_state['characters'] = characters; logger.info(f"Loaded {len(characters)} characters.") + else: st.session_state['characters'] = []; logger.warning("characters.json is not a list, resetting."); os.remove("characters.json") except (json.JSONDecodeError, IOError) as e: logger.error(f"Failed to load or decode characters.json: {e}") st.error(f"Error loading character file: {e}. Starting fresh.") st.session_state['characters'] = [] try: corrupt_filename = f"characters_corrupt_{int(time.time())}.json" - shutil.copy("characters.json", corrupt_filename) - logger.info(f"Backed up corrupted character file to {corrupt_filename}") - os.remove("characters.json") - except Exception as backup_e: - logger.error(f"Could not backup corrupted character file: {backup_e}") + shutil.copy("characters.json", corrupt_filename); logger.info(f"Backed up corrupted character file to {corrupt_filename}"); os.remove("characters.json") + except Exception as backup_e: logger.error(f"Could not backup corrupted character file: {backup_e}") + +# --- Utility: Clean stems (Keep from previous) ---------- def clean_stem(fn: str) -> str: name = os.path.splitext(os.path.basename(fn))[0] name = name.replace('-', ' ').replace('_', ' ') return name.strip().title() -def make_image_sized_pdf(sources, is_markdown_flags): - if not sources: - st.warning("No sources provided for PDF generation.") - return None - buf = BytesIO() - styles = getSampleStyleSheet() - md_style = ParagraphStyle( - name='Markdown', - fontSize=10, - leading=12, - spaceAfter=6, - alignment=TA_JUSTIFY, - fontName='Helvetica' - ) - doc = SimpleDocTemplate(buf, pagesize=letter, rightMargin=36, leftMargin=36, topMargin=36, bottomMargin=36) - story = [] + +# --- PDF Creation: Image Sized + Captions (Keep from previous) --- +def make_image_sized_pdf(sources): + if not sources: st.warning("No image sources provided for PDF generation."); return None + buf = io.BytesIO() + c = canvas.Canvas(buf, pagesize=letter) # Default letter try: - for idx, (src, is_md) in enumerate(zip(sources, is_markdown_flags), start=1): + for idx, src in enumerate(sources, start=1): status_placeholder = st.empty() - filename = 'page_' + str(idx) status_placeholder.info(f"Adding page {idx}/{len(sources)}: {os.path.basename(str(src))}...") try: - if is_md: - with open(src, 'r', encoding='utf-8') as f: - content = f.read() - content = re.sub(r'!\[.*?\]\(.*?\)', '', content) - paragraphs = content.split('\n\n') - for para in paragraphs: - if para.strip(): - story.append(Paragraph(para.strip(), md_style)) - story.append(PageBreak()) - status_placeholder.success(f"Added markdown page {idx}/{len(sources)}: {filename}") + 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: - if isinstance(src, str): - if not os.path.exists(src): - logger.warning(f"Image file not found: {src}. Skipping.") - status_placeholder.warning(f"Skipping missing file: {os.path.basename(src)}") - continue - img_obj = Image.open(src) - filename = os.path.basename(src) - else: - src.seek(0) - img_obj = Image.open(src) - filename = getattr(src, 'name', f'uploaded_image_{idx}') - src.seek(0) - with img_obj: - iw, ih = img_obj.size - if iw <= 0 or ih <= 0: - logger.warning(f"Invalid image dimensions ({iw}x{ih}) for {filename}. Skipping.") - status_placeholder.warning(f"Skipping invalid image: {filename}") - continue - cap_h = 30 - c = canvas.Canvas(BytesIO(), pagesize=(iw, ih + cap_h)) - img_reader = ImageReader(img_obj) - c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto') - caption = clean_stem(filename) - c.setFont('Helvetica', 12) - c.setFillColorRGB(0, 0, 0) - c.drawCentredString(iw / 2, cap_h / 2 + 3, caption) - c.setFont('Helvetica', 8) - c.setFillColorRGB(0.5, 0.5, 0.5) - c.drawRightString(iw - 10, 8, f"Page {idx}") - c.save() - story.append(PageBreak()) - status_placeholder.success(f"Added image page {idx}/{len(sources)}: {filename}") - except Exception as e: - logger.error(f"Error processing source {src}: {e}") - status_placeholder.error(f"Error adding page {idx}: {e}") - doc.build(story) - buf.seek(0) - if buf.getbuffer().nbytes < 100: - st.error("PDF generation resulted in an empty file.") - return None + src.seek(0); img_obj = Image.open(src); filename = getattr(src, 'name', f'uploaded_image_{idx}'); src.seek(0) + + with img_obj: + iw, ih = img_obj.size + if iw <= 0 or ih <= 0: logger.warning(f"Invalid image dimensions ({iw}x{ih}) for {filename}. Skipping."); status_placeholder.warning(f"Skipping invalid image: {filename}"); continue + cap_h = 30; pw, ph = iw, ih + cap_h + c.setPageSize((pw, ph)) + img_reader = ImageReader(img_obj) + c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto') + caption = clean_stem(filename); c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, caption) + c.setFont('Helvetica', 8); c.setFillColorRGB(0.5, 0.5, 0.5); c.drawRightString(pw - 10, 8, f"Page {idx}") + c.showPage() + status_placeholder.success(f"Added page {idx}/{len(sources)}: {filename}") + + except (IOError, OSError, UnidentifiedImageError) as img_err: logger.error(f"Error processing image {src}: {img_err}"); status_placeholder.error(f"Error adding page {idx}: {img_err}") + except Exception as e: logger.error(f"Unexpected error adding page {idx} ({src}): {e}"); status_placeholder.error(f"Unexpected error on page {idx}: {e}") + + c.save(); buf.seek(0) + if buf.getbuffer().nbytes < 100: st.error("PDF generation resulted in an empty file."); return None return buf.getvalue() except Exception as e: logger.error(f"Fatal error during PDF generation: {e}") st.error(f"PDF Generation Failed: {e}") return None -def update_gallery(gallery_type='image'): - container = st.session_state['asset_gallery_container'][gallery_type] - with container: - st.markdown(f"### {gallery_type.capitalize()} Gallery πŸ“Έ") - files = get_typed_gallery_files(gallery_type) - if not files: - st.info(f"No {gallery_type} assets found yet.") - return - st.caption(f"Found {len(files)} assets:") - for idx, file in enumerate(files[:st.session_state.gallery_size]): - st.session_state['unique_counter'] += 1 - unique_id = st.session_state['unique_counter'] - item_key_base = f"{gallery_type}_gallery_item_{os.path.basename(file)}_{unique_id}" - basename = os.path.basename(file) - st.markdown(f"**{basename}**") - try: - file_ext = os.path.splitext(file)[1].lower() - if gallery_type == 'image' and file_ext in ['.png', '.jpg', '.jpeg']: - with st.expander("Preview", expanded=False): - st.image(Image.open(file), use_container_width=True) - elif gallery_type == 'pdf' and file_ext == '.pdf': - with st.expander("Preview (Page 1)", expanded=False): - doc = fitz.open(file) - if len(doc) > 0: - pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) - img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) - st.image(img, use_container_width=True) - else: - st.warning("Empty PDF") - doc.close() - elif gallery_type == 'md' and file_ext == '.md': - with st.expander("Preview (Start)", expanded=False): - with open(file, 'r', encoding='utf-8', errors='ignore') as f: - content_preview = f.read(200) - st.code(content_preview + "...", language='markdown') - action_cols = st.columns(3) - with action_cols[0]: - checkbox_key = f"cb_{item_key_base}" - st.session_state['asset_checkboxes'][gallery_type][file] = st.checkbox( - "Select", - value=st.session_state['asset_checkboxes'][gallery_type].get(file, False), - key=checkbox_key - ) - with action_cols[1]: - mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.md': 'text/markdown'} - mime_type = mime_map.get(file_ext, "application/octet-stream") - dl_key = f"dl_{item_key_base}" + +# --- Sidebar Gallery Update Function (MODIFIED) -------- +def update_gallery(): + st.sidebar.markdown("### Asset Gallery πŸ“ΈπŸ“–") + + all_files = get_gallery_files() # Get currently available files + + if not all_files: + st.sidebar.info("No assets (images, PDFs, text files) found yet.") + return + + st.sidebar.caption(f"Found {len(all_files)} assets:") + + for idx, file in enumerate(all_files): + st.session_state['unique_counter'] += 1 + unique_id = st.session_state['unique_counter'] + item_key_base = f"gallery_item_{os.path.basename(file)}_{unique_id}" + basename = os.path.basename(file) + st.sidebar.markdown(f"**{basename}**") # Display filename clearly + + try: + file_ext = os.path.splitext(file)[1].lower() + # Display previews + if file_ext in ['.png', '.jpg', '.jpeg']: + # Add expander for large galleries + with st.sidebar.expander("Preview", expanded=False): + st.image(Image.open(file), use_container_width=True) + elif file_ext == '.pdf': + with st.sidebar.expander("Preview (Page 1)", expanded=False): + doc = fitz.open(file) + if len(doc) > 0: + pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) # Smaller preview + img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) + st.image(img, use_container_width=True) + else: + st.warning("Empty PDF") + doc.close() + elif file_ext in ['.md', '.txt']: + with st.sidebar.expander("Preview (Start)", expanded=False): + with open(file, 'r', encoding='utf-8', errors='ignore') as f: + content_preview = f.read(200) # Show first 200 chars + st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text') + + # --- Actions for the file (Select, Download, Delete) --- + action_cols = st.sidebar.columns(3) # Use columns for buttons + with action_cols[0]: + checkbox_key = f"cb_{item_key_base}" + st.session_state['asset_checkboxes'][file] = st.checkbox( + "Select", + value=st.session_state['asset_checkboxes'].get(file, False), + key=checkbox_key + ) + with action_cols[1]: + mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'} + mime_type = mime_map.get(file_ext, "application/octet-stream") + # Use button for download to avoid complex HTML link generation issues sometimes + dl_key = f"dl_{item_key_base}" + try: + with open(file, "rb") as fp: + st.download_button( + label="πŸ“₯", + data=fp, + file_name=basename, + mime=mime_type, + key=dl_key, + help="Download this file" + ) + except Exception as dl_e: + st.error(f"DL Err: {dl_e}") + + with action_cols[2]: + delete_key = f"del_{item_key_base}" + if st.button("πŸ—‘οΈ", key=delete_key, help=f"Delete {basename}"): try: - with open(file, "rb") as fp: - st.download_button( - label="πŸ“₯", - data=fp, - file_name=basename, - mime=mime_type, - key=dl_key, - help="Download this file" - ) - except Exception as dl_e: - st.error(f"Download Error: {dl_e}") - with action_cols[2]: - delete_key = f"del_{item_key_base}" - if st.button("πŸ—‘οΈ", key=delete_key, help=f"Delete {basename}"): - try: - os.remove(file) - st.session_state['asset_checkboxes'][gallery_type].pop(file, None) - if file in st.session_state.get('layout_snapshots', []): - st.session_state['layout_snapshots'].remove(file) - logger.info(f"Deleted {gallery_type} asset: {file}") - st.toast(f"Deleted {basename}!", icon="βœ…") - st.rerun() - except OSError as e: - logger.error(f"Error deleting file {file}: {e}") - st.error(f"Could not delete {basename}") - except Exception as e: - st.error(f"Error displaying {basename}: {e}") - logger.error(f"Error displaying asset {file}: {e}") - st.markdown("---") - -# --- UI Elements --- -st.sidebar.subheader("πŸ€– AI Settings") + os.remove(file) + st.session_state['asset_checkboxes'].pop(file, None) # Remove from selection state + # Remove from layout_snapshots if present + if file in st.session_state.get('layout_snapshots', []): + st.session_state['layout_snapshots'].remove(file) + logger.info(f"Deleted asset: {file}") + st.toast(f"Deleted {basename}!", icon="βœ…") # Use toast for less intrusive feedback + # REMOVED st.rerun() - Rely on file watcher + except OSError as e: + logger.error(f"Error deleting file {file}: {e}") + st.error(f"Could not delete {basename}") + # Trigger a rerun MANUALLY after deletion completes if file watcher is unreliable + st.rerun() + + + except (fitz.fitz.FileNotFoundError, FileNotFoundError): + st.sidebar.error(f"File not found: {basename}") + st.session_state['asset_checkboxes'].pop(file, None) # Clean up state + except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: + st.sidebar.error(f"Corrupt PDF: {basename}") + logger.warning(f"Error opening PDF {file}: {pdf_err}") + except UnidentifiedImageError: + st.sidebar.error(f"Invalid Image: {basename}") + logger.warning(f"Cannot identify image file {file}") + except Exception as e: + st.sidebar.error(f"Error: {basename}") + logger.error(f"Error displaying asset {file}: {e}") + + st.sidebar.markdown("---") # Separator between items + +# --- UI Elements ----------------------------------------- + +# --- Sidebar: HF Inference Settings --- +st.sidebar.subheader("πŸ€– Hugging Face Settings") +st.sidebar.markdown("Configure API inference or select local models.") + +# API Settings Expander with st.sidebar.expander("API Inference Settings", expanded=False): st.session_state.hf_custom_key = st.text_input( - "Custom HF Token", + "Custom HF Token (BYOK)", value=st.session_state.get('hf_custom_key', ""), type="password", - key="hf_custom_key_input" + 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"HF Token Status: {token_status}") + 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( - "HF Inference Provider", + "Inference Provider", options=providers_list, index=providers_list.index(st.session_state.get('hf_provider', DEFAULT_PROVIDER)), - key="hf_provider_select" + 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 HF API Model ID", + "Custom API Model ID", value=st.session_state.get('hf_custom_api_model', ""), - key="hf_custom_model_input" + key="hf_custom_model_input", + placeholder="e.g., google/gemma-2-9b-it", + help="Overrides the featured model selection below if provided." ) - effective_hf_model = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model + # 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 HF API Model", + "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" + key="hf_featured_model_select", + help="Select a common model. Ignored if Custom API Model ID is set." ) - st.caption(f"Effective HF API Model: {effective_hf_model}") - if _openai_available: - st.session_state.openai_selected_model = st.selectbox( - "OpenAI Model", - options=OPENAI_MODELS_LIST, - index=OPENAI_MODELS_LIST.index(st.session_state.get('openai_selected_model', OPENAI_MODELS_LIST[0])), - key="openai_model_select" - ) + st.caption(f"Effective API Model: {effective_api_model}") + +# Local Model Selection Expander with st.sidebar.expander("Local Model Selection", expanded=True): if not _transformers_available: - st.warning("Transformers library not found. Cannot load local models.") + 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', "None") + 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" + 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" + 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'}") + 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.") + 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), key="param_max_tokens") + 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") - st.session_state.gen_frequency_penalty = st.slider("Repetition Penalty", 0.0, 1.0, st.session_state.get('gen_frequency_penalty', 0.0), step=0.05, key="param_repetition") - st.session_state.gen_seed = st.slider("Seed", -1, 65535, st.session_state.get('gen_seed', -1), step=1, key="param_seed") + # Note: HF often uses repetition_penalty instead of frequency_penalty. We'll use it here. + st.session_state.gen_frequency_penalty = st.slider("Repetition Penalty", 1.0, 2.0, st.session_state.get('gen_frequency_penalty', 0.0)+1.0, step=0.05, key="param_repetition", help="1.0 means no penalty.") + st.session_state.gen_seed = st.slider("Seed", -1, 65535, st.session_state.get('gen_seed', -1), step=1, key="param_seed", help="-1 for random.") + + +st.sidebar.markdown("---") # Separator before gallery settings + +# --- ADDED: Gallery Settings Section --- st.sidebar.subheader("πŸ–ΌοΈ Gallery Settings") st.slider( "Max Items Shown", min_value=2, - max_value=50, + max_value=50, # Adjust max if needed value=st.session_state.get('gallery_size', 10), - key="gallery_size_slider" + key="gallery_size_slider", # Keep the key, define it ONCE here + help="Controls the maximum number of assets displayed in the sidebar gallery." ) -st.session_state.gallery_size = st.session_state.gallery_size_slider -st.sidebar.markdown("---") -update_gallery('image') -update_gallery('md') -update_gallery('pdf') - -# --- Main Application --- -st.title("Vision & Layout Titans πŸš€πŸ–ΌοΈπŸ“„") -st.markdown("Create PDFs from images and markdown, process with AI, and manage characters.") -tabs = st.tabs([ - "Image/MD->PDF Layout πŸ–ΌοΈβž‘οΈπŸ“„", +st.session_state.gallery_size = st.session_state.gallery_size_slider # Ensure sync +st.sidebar.markdown("---") # Separator after gallery settings + + + + +# --- App Title ------------------------------------------- +st.title("Vision & Layout Titans (HF) πŸš€πŸ–ΌοΈπŸ“„") +st.markdown("Combined App: Image-to-PDF Layout + Hugging Face Powered AI Tools") + +# --- Main Application Tabs ------------------------------- +tab_list = [ + "Image->PDF Layout πŸ–ΌοΈβž‘οΈπŸ“„", # From App 1 + "Camera Snap πŸ“·", + "Download PDFs πŸ“₯", + "Build Titan (Local Models) 🌱", # Renamed for clarity + "Text Process (HF) πŸ“", # New tab for text + "Image Process (HF) πŸ–ΌοΈ", # New tab for image + "Test OCR (HF) πŸ”", # Renamed + "Character Editor πŸ§‘β€πŸŽ¨", + "Character Gallery πŸ–ΌοΈ", + # Original Tabs (potentially redundant or integrated now): + # "PDF Process πŸ“„", (Integrated into Text/Image process conceptually) + # "MD Gallery & Process πŸ“š", (Use Text Process tab) + # "Test Image Gen 🎨", (Separate Diffusion logic) +] +# Filter out redundant tabs if they are fully replaced +# Example: If MD Gallery is fully handled by Text Process, remove it. For now, keep most. +# Let's keep PDF Process and Image Process separate for clarity of input type, but use the new HF functions +tabs_to_create = [ + "Image->PDF Layout πŸ–ΌοΈβž‘οΈπŸ“„", "Camera Snap πŸ“·", "Download PDFs πŸ“₯", "Build Titan (Local Models) 🌱", - "PDF Process (AI) πŸ“„", - "Image Process (AI) πŸ–ΌοΈ", - "Text Process (AI) πŸ“", - "Test OCR (AI) πŸ”", - "Test Image Gen (Diffusers) 🎨", + "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(tabs_to_create) + +# --- Tab Implementations --- +# --- Tab 1: Image -> PDF Layout (Keep from previous merge) --- with tabs[0]: - st.header("Image/Markdown to PDF Layout Generator") - st.markdown("Select images and markdown files, reorder them, and generate a PDF.") + # ... (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. Select Assets") - selected_images = [f for f in get_typed_gallery_files('image') if st.session_state['asset_checkboxes']['image'].get(f, False)] - selected_mds = [f for f in get_typed_gallery_files('md') if st.session_state['asset_checkboxes']['md'].get(f, False)] - st.write(f"Selected Images: {len(selected_images)}") - st.write(f"Selected Markdown Files: {len(selected_mds)}") + st.subheader("A. Scan or Upload Images") + layout_cam = st.camera_input("πŸ“Έ Scan Document for Layout PDF", key="layout_cam") + if layout_cam: + now = datetime.now(pytz.timezone("US/Central")) + scan_name = generate_filename(f"layout_scan_{now.strftime('%a').upper()}", "png") + try: + with open(scan_name, "wb") as f: f.write(layout_cam.getvalue()) + st.image(Image.open(scan_name), caption=f"Scanned: {scan_name}", use_container_width=True) + if scan_name not in st.session_state['layout_snapshots']: st.session_state['layout_snapshots'].append(scan_name) + st.success(f"Scan saved as {scan_name}") + update_gallery(); # Add to gallery + except Exception as e: st.error(f"Failed to save scan: {e}"); logger.error(f"Failed to save camera scan {scan_name}: {e}") + + layout_uploads = st.file_uploader("πŸ“‚ Upload PNG/JPG Images for Layout PDF", type=["png","jpg","jpeg"], accept_multiple_files=True, key="layout_uploader") + if layout_uploads: st.session_state['layout_new_uploads'] = layout_uploads # Store for processing below with col2: st.subheader("B. Review and Reorder") layout_records = [] - for idx, path in enumerate(selected_images + selected_mds, start=1): - is_md = path in selected_mds + processed_snapshots = set() + # Process snapshots + for idx, path in enumerate(st.session_state.get('layout_snapshots', [])): + if path not in processed_snapshots and os.path.exists(path): + try: + with Image.open(path) as im: w, h = im.size; ar = round(w / h, 2) if h > 0 else 0; orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait") + layout_records.append({"filename": os.path.basename(path), "source": path, "width": w, "height": h, "aspect_ratio": ar, "orientation": orient, "order": idx, "type": "Scan"}) + processed_snapshots.add(path) + except Exception as e: logger.warning(f"Could not process snapshot {path}: {e}"); st.warning(f"Skipping invalid snapshot: {os.path.basename(path)}") + # Process current uploads + current_uploads = st.session_state.get('layout_new_uploads', []) + if current_uploads: + start_idx = len(layout_records) + for jdx, f_obj in enumerate(current_uploads, start=start_idx): + try: + f_obj.seek(0) + with Image.open(f_obj) as im: w, h = im.size; ar = round(w / h, 2) if h > 0 else 0; orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait") + layout_records.append({"filename": f_obj.name, "source": f_obj, "width": w, "height": h, "aspect_ratio": ar, "orientation": orient, "order": jdx, "type": "Upload"}) + f_obj.seek(0) + except Exception as e: logger.warning(f"Could not process uploaded file {f_obj.name}: {e}"); st.warning(f"Skipping invalid upload: {f_obj.name}") + + if not layout_records: st.info("Scan or upload images using the controls on the left.") + else: + layout_df = pd.DataFrame(layout_records); dims = st.multiselect("Include orientations:", options=["Landscape","Portrait","Square"], default=["Landscape","Portrait","Square"], key="layout_dims_filter") + filtered_df = layout_df[layout_df['orientation'].isin(dims)].copy() if dims else layout_df.copy() + filtered_df['order'] = filtered_df['order'].astype(int); filtered_df = filtered_df.sort_values('order').reset_index(drop=True) + st.markdown("Edit 'Order' column or drag rows to set PDF page sequence:") + edited_df = st.data_editor(filtered_df, column_config={"filename": st.column_config.TextColumn("Filename", disabled=True), "source": None, "width": st.column_config.NumberColumn("Width", disabled=True), "height": st.column_config.NumberColumn("Height", disabled=True), "aspect_ratio": st.column_config.NumberColumn("Aspect Ratio", format="%.2f", disabled=True), "orientation": st.column_config.TextColumn("Orientation", disabled=True), "type": st.column_config.TextColumn("Source Type", disabled=True), "order": st.column_config.NumberColumn("Order", min_value=0, step=1, required=True)}, hide_index=True, use_container_width=True, num_rows="dynamic", key="layout_editor") + ordered_layout_df = edited_df.sort_values('order').reset_index(drop=True) + ordered_sources_for_pdf = ordered_layout_df['source'].tolist() + + st.subheader("C. Generate & Download PDF") + if st.button("πŸ–‹οΈ Generate Image-Sized PDF", key="generate_layout_pdf"): + if not ordered_sources_for_pdf: st.warning("No images selected or available after filtering.") + else: + with st.spinner("Generating PDF..."): pdf_bytes = make_image_sized_pdf(ordered_sources_for_pdf) + if pdf_bytes: + now = datetime.now(pytz.timezone("US/Central")); prefix = now.strftime("%Y%m%d-%H%M%p") + stems = [clean_stem(s) if isinstance(s, str) else clean_stem(getattr(s, 'name', 'upload')) for s in ordered_sources_for_pdf[:4]] + basename = " - ".join(stems) or "Layout"; pdf_fname = f"{prefix}_{basename}.pdf"; pdf_fname = re.sub(r'[^\w\- \.]', '_', pdf_fname) + st.success(f"βœ… PDF ready: **{pdf_fname}**") + st.download_button("⬇️ Download PDF", data=pdf_bytes, file_name=pdf_fname, mime="application/pdf", key="download_layout_pdf") + st.markdown("#### Preview First Page") + try: + doc = fitz.open(stream=pdf_bytes, filetype='pdf') + if len(doc) > 0: pix = doc[0].get_pixmap(matrix=fitz.Matrix(1.0, 1.0)); preview_img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); st.image(preview_img, caption=f"Preview of {pdf_fname} (Page 1)", use_container_width=True) + else: st.warning("Generated PDF appears empty.") + doc.close() + except Exception as preview_err: st.warning(f"Could not generate PDF preview: {preview_err}"); logger.warning(f"PDF preview error for {pdf_fname}: {preview_err}") + else: st.error("PDF generation failed. Check logs or image files.") + + +# --- Tab 2: Camera Snap (Keep from previous merge) --- +with tabs[1]: + # ... (Code from previous merge for this tab) ... + st.header("Camera Snap πŸ“·") + st.subheader("Single Capture (Adds to General Gallery)") + cols = st.columns(2) + with cols[0]: + cam0_img = st.camera_input("Take a picture - Cam 0", key="main_cam0") + if cam0_img: + filename = generate_filename("cam0_snap"); + if st.session_state.get('cam0_file') and os.path.exists(st.session_state['cam0_file']): + try: + os.remove(st.session_state['cam0_file']) + except OSError: + pass try: - if is_md: - with open(path, 'r', encoding='utf-8') as f: - content = f.read(50) - layout_records.append({ - "filename": os.path.basename(path), - "source": path, - "type": "Markdown", - "preview": content + "...", - "order": idx - }) + with open(filename, "wb") as f: f.write(cam0_img.getvalue()) + st.session_state['cam0_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 0: {filename}"); st.image(Image.open(filename), caption="Camera 0 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 0: {filename}"); st.success(f"Saved {filename}") + update_gallery(); + except Exception as e: + st.error(f"Failed to save Cam 0 snap: {e}"); logger.error(f"Failed to save Cam 0 snap {filename}: {e}") + with cols[1]: + cam1_img = st.camera_input("Take a picture - Cam 1", key="main_cam1") + if cam1_img: + filename = generate_filename("cam1_snap") + if st.session_state.get('cam1_file') and os.path.exists(st.session_state['cam1_file']): + try: + os.remove(st.session_state['cam1_file']) + except OSError: + pass + try: + with open(filename, "wb") as f: f.write(cam1_img.getvalue()) + st.session_state['cam1_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 1: {filename}"); st.image(Image.open(filename), caption="Camera 1 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 1: {filename}"); st.success(f"Saved {filename}") + update_gallery(); + except Exception as e: st.error(f"Failed to save Cam 1 snap: {e}"); logger.error(f"Failed to save Cam 1 snap {filename}: {e}") + + +# --- Tab 3: Download PDFs (Keep from previous merge) --- +with tabs[2]: + # ... (Code from previous merge for this tab) ... + st.header("Download PDFs πŸ“₯") + st.markdown("Download PDFs from URLs and optionally create image snapshots.") + if st.button("Load Example arXiv URLs πŸ“š", key="load_examples"): + example_urls = ["https://arxiv.org/pdf/2308.03892", "https://arxiv.org/pdf/1706.03762", "https://arxiv.org/pdf/2402.17764", "https://www.clickdimensions.com/links/ACCERL/"] + st.session_state['pdf_urls_input'] = "\n".join(example_urls) + url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls_input', ""), height=150, key="pdf_urls_textarea") + if st.button("Robo-Download PDFs πŸ€–", key="download_pdfs_button"): + urls = [url.strip() for url in url_input.strip().split("\n") if url.strip()] + if not urls: st.warning("Please enter at least one URL.") + else: + progress_bar = st.progress(0); status_text = st.empty(); total_urls = len(urls); download_count = 0; existing_pdfs = get_pdf_files() + for idx, url in enumerate(urls): + output_path = pdf_url_to_filename(url); status_text.text(f"Processing {idx + 1}/{total_urls}: {os.path.basename(output_path)}..."); progress_bar.progress((idx + 1) / total_urls) + if output_path in existing_pdfs: st.info(f"Already exists: {os.path.basename(output_path)}"); st.session_state['downloaded_pdfs'][url] = output_path; st.session_state['asset_checkboxes'][output_path] = st.session_state['asset_checkboxes'].get(output_path, False) else: - with Image.open(path) as im: - w, h = im.size - ar = round(w / h, 2) if h > 0 else 0 - orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait") - layout_records.append({ - "filename": os.path.basename(path), - "source": path, - "type": "Image", - "width": w, - "height": h, - "aspect_ratio": ar, - "orientation": orient, - "order": idx - }) - except Exception as e: - logger.warning(f"Could not process {path}: {e}") - st.warning(f"Skipping invalid file: {os.path.basename(path)}") - if not layout_records: - st.infoperiod \ No newline at end of file + if download_pdf(url, output_path): st.session_state['downloaded_pdfs'][url] = output_path; logger.info(f"Downloaded PDF from {url} to {output_path}"); st.session_state['history'].append(f"Downloaded PDF: {output_path}"); st.session_state['asset_checkboxes'][output_path] = False; download_count += 1; existing_pdfs.append(output_path) + else: st.error(f"Failed to download: {url}") + status_text.success(f"Download process complete! Successfully downloaded {download_count} new PDFs.") + if download_count > 0: update_gallery(); + + st.subheader("Create Snapshots from Gallery PDFs") + snapshot_mode = st.selectbox("Snapshot Mode", ["First Page (High-Res)", "First Two Pages (High-Res)", "All Pages (High-Res)", "First Page (Low-Res Preview)"], key="pdf_snapshot_mode") + resolution_map = {"First Page (High-Res)": 2.0, "First Two Pages (High-Res)": 2.0, "All Pages (High-Res)": 2.0, "First Page (Low-Res Preview)": 1.0} + mode_key_map = {"First Page (High-Res)": "single", "First Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages", "First Page (Low-Res Preview)": "single"} + resolution = resolution_map[snapshot_mode]; mode_key = mode_key_map[snapshot_mode] + if st.button("Snapshot Selected PDFs πŸ“Έ", key="snapshot_selected_pdfs"): + selected_pdfs = [path for path in get_gallery_files(['pdf']) if st.session_state['asset_checkboxes'].get(path, False)] + if not selected_pdfs: st.warning("No PDFs selected in the sidebar gallery!") + else: + st.info(f"Starting snapshot process for {len(selected_pdfs)} selected PDF(s)..."); snapshot_count = 0; total_snapshots_generated = 0 + for pdf_path in selected_pdfs: + if not os.path.exists(pdf_path): st.warning(f"File not found: {pdf_path}. Skipping."); continue + new_snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key, resolution)) + if new_snapshots: + snapshot_count += 1; total_snapshots_generated += len(new_snapshots) + st.write(f"Snapshots for {os.path.basename(pdf_path)}:"); cols = st.columns(3) + for i, snap_path in enumerate(new_snapshots): + with cols[i % 3]: st.image(Image.open(snap_path), caption=os.path.basename(snap_path), use_container_width=True); st.session_state['asset_checkboxes'][snap_path] = False # Add to gallery + if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery(); + else: st.warning("No snapshots were generated. Check logs or PDF files.") + +# --- Tab 4: Build Titan (Local Models) --- +with tabs[3]: + st.header("Build Titan (Local Models) 🌱") + st.markdown("Download and save models from Hugging Face Hub for local use.") + + if not _transformers_available: + st.error("Transformers library not available. Cannot download or load local models.") + else: + build_model_type = st.selectbox( + "Select Model Type", + ["Causal LM", "Vision/Multimodal", "OCR", "Diffusion"], # Added more types + key="build_type_local" + ) + + st.subheader(f"Download {build_model_type} Model") + # Model ID Input (allow searching/pasting) + hf_model_id = st.text_input( + "Hugging Face Model ID", + placeholder=f"e.g., {'google/gemma-2-9b-it' if build_model_type == 'Causal LM' else 'llava-hf/llava-1.5-7b-hf' if build_model_type == 'Vision/Multimodal' else 'microsoft/trocr-base-handwritten' if build_model_type == 'OCR' else 'stabilityai/stable-diffusion-xl-base-1.0'}", + key="build_hf_model_id" + ) + local_model_name = st.text_input( + "Local Name for this Model", + value=f"{build_model_type.split('/')[0].lower()}_{os.path.basename(hf_model_id).replace('.','') if hf_model_id else 'model'}", + key="build_local_name", + help="A unique name to identify this model locally." + ) + + # Add a note about token requirements for gated models + st.info("Private or gated models require a valid Hugging Face token (set via HF_TOKEN env var or the Custom Key in sidebar API settings).") + + if st.button(f"Download & Save '{hf_model_id}' Locally", key="build_download_button", disabled=not hf_model_id or not local_model_name): + # Validate local name uniqueness + if local_model_name in [os.path.basename(p) for p in st.session_state.get('local_models', {})]: + st.error(f"A local model named '{local_model_name}' already exists. Choose a different name.") + else: + model_type_map = { + "Causal LM": "causal", "Vision/Multimodal": "vision", "OCR": "ocr", "Diffusion": "diffusion" + } + model_type_short = model_type_map.get(build_model_type, "unknown") + + config = LocalModelConfig( + name=local_model_name, + hf_id=hf_model_id, + model_type=model_type_short + ) + save_path = config.get_full_path() + os.makedirs(os.path.dirname(save_path), exist_ok=True) + + st.info(f"Attempting to download '{hf_model_id}' to '{save_path}'...") + progress_bar_build = st.progress(0) + status_text_build = st.empty() + token_build = st.session_state.hf_custom_key or HF_TOKEN or None + + try: + if build_model_type == "Diffusion": + # Use Diffusers library download + if not _diffusers_available: raise ImportError("Diffusers library required for diffusion models.") + # Diffusers downloads directly, no explicit save needed after load typically + status_text_build.text("Downloading diffusion model pipeline...") + pipeline_obj = StableDiffusionPipeline.from_pretrained(hf_model_id, token=token_build) + status_text_build.text("Saving diffusion model pipeline...") + pipeline_obj.save_pretrained(save_path) + # Store info, but maybe not the full pipeline object in session state due to size + st.session_state.local_models[save_path] = {'type': 'diffusion', 'hf_id': hf_model_id, 'model':None, 'tokenizer':None} # Mark as downloaded + st.success(f"Diffusion model '{hf_model_id}' downloaded and saved to {save_path}") + + else: + # Use Transformers library download + status_text_build.text("Downloading model components...") + # Determine AutoModel class based on type (can be refined) + if model_type_short == 'causal': + model_class = AutoModelForCausalLM + tokenizer_class = AutoTokenizer + processor_class = None + elif model_type_short == 'vision': + model_class = AutoModelForVision2Seq # Common for VQA/Captioning + processor_class = AutoProcessor # Handles image+text + tokenizer_class = None # Usually part of processor + elif model_type_short == 'ocr': + model_class = AutoModelForVision2Seq # TrOCR uses this + processor_class = AutoProcessor + tokenizer_class = None + else: + raise ValueError(f"Unknown model type for downloading: {model_type_short}") + + # Download and save model + model_obj = model_class.from_pretrained(hf_model_id, token=token_build) + model_obj.save_pretrained(save_path) + status_text_build.text(f"Model saved. Downloading processor/tokenizer...") + + # Download and save tokenizer/processor + if processor_class: + processor_obj = processor_class.from_pretrained(hf_model_id, token=token_build) + processor_obj.save_pretrained(save_path) + tokenizer_obj = getattr(processor_obj, 'tokenizer', None) # Get tokenizer from processor if exists + elif tokenizer_class: + tokenizer_obj = tokenizer_class.from_pretrained(hf_model_id, token=token_build) + tokenizer_obj.save_pretrained(save_path) + processor_obj = None # No separate processor + else: # Should not happen with current logic + tokenizer_obj = None + processor_obj = None + + # --- Load into memory and store in session state --- + # This might consume significant memory! Consider loading on demand instead. + status_text_build.text(f"Loading '{local_model_name}' into memory...") + device = "cuda" if torch.cuda.is_available() else "cpu" + + reloaded_model = model_class.from_pretrained(save_path).to(device) + reloaded_processor = processor_class.from_pretrained(save_path) if processor_class else None + reloaded_tokenizer = tokenizer_class.from_pretrained(save_path) if tokenizer_class and not reloaded_processor else getattr(reloaded_processor, 'tokenizer', None) + + st.session_state.local_models[save_path] = { + 'type': model_type_short, + 'hf_id': hf_model_id, + 'model': reloaded_model, + 'tokenizer': reloaded_tokenizer, + 'processor': reloaded_processor, # Store processor if it exists + } + st.success(f"{build_model_type} model '{hf_model_id}' downloaded to {save_path} and loaded into memory ({device}).") + # Optionally select the newly loaded model + st.session_state.selected_local_model_path = save_path + + + except (RepositoryNotFoundError, GatedRepoError) as e: + st.error(f"Download failed: Repository not found or requires specific access/token. Check Model ID and your HF token. Error: {e}") + logger.error(f"Download failed for {hf_model_id}: {e}") + if os.path.exists(save_path): shutil.rmtree(save_path) # Clean up partial download + except ImportError as e: + st.error(f"Download failed: Required library missing. {e}") + logger.error(f"ImportError during download of {hf_model_id}: {e}") + except Exception as e: + st.error(f"An unexpected error occurred during download: {e}") + logger.error(f"Download failed for {hf_model_id}: {e}") + if os.path.exists(save_path): shutil.rmtree(save_path) # Clean up + + finally: + progress_bar_build.progress(1.0) + status_text_build.empty() + + st.subheader("Manage Local Models") + loaded_model_paths = list(st.session_state.get('local_models', {}).keys()) + if not loaded_model_paths: + st.info("No models downloaded yet.") + else: + models_df_data = [] + for path, data in st.session_state.local_models.items(): + models_df_data.append({ + "Local Name": os.path.basename(path), + "Type": data.get('type', 'N/A'), + "HF ID": data.get('hf_id', 'N/A'), + "Loaded": "Yes" if data.get('model') else "No (Info only)", + "Path": path + }) + models_df = pd.DataFrame(models_df_data) + st.dataframe(models_df, use_container_width=True, hide_index=True, column_order=["Local Name", "Type", "HF ID", "Loaded"]) + + model_to_delete = st.selectbox("Select model to delete", [""] + [os.path.basename(p) for p in loaded_model_paths], key="delete_model_select") + if model_to_delete and st.button(f"Delete Local Model '{model_to_delete}'", type="primary"): + path_to_delete = next((p for p in loaded_model_paths if os.path.basename(p) == model_to_delete), None) + if path_to_delete: + try: + # Remove from session state first + del st.session_state.local_models[path_to_delete] + if st.session_state.selected_local_model_path == path_to_delete: + st.session_state.selected_local_model_path = None + # Delete from disk + if os.path.exists(path_to_delete): + shutil.rmtree(path_to_delete) + st.success(f"Deleted model '{model_to_delete}' and its files.") + logger.info(f"Deleted local model: {path_to_delete}") + except Exception as e: + st.error(f"Failed to delete model '{model_to_delete}': {e}") + logger.error(f"Failed to delete model {path_to_delete}: {e}") + + +# --- Tab 5: PDF Process (HF) --- +with tabs[4]: + st.header("PDF Process with HF Models πŸ“„") + st.markdown("Upload PDFs, view pages, and extract text using selected HF models (API or Local).") + + # Inference Source Selection + pdf_use_api = st.radio( + "Choose Processing Method", + ["Hugging Face API", "Loaded Local Model"], + key="pdf_process_source", + horizontal=True, + help="API uses settings from sidebar. Local uses the selected local model (if suitable)." + ) + + if pdf_use_api == "Hugging Face API": + st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}") + else: + if st.session_state.selected_local_model_path: + st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}") + else: + st.warning("No local model selected. Please select one in the sidebar.") + + uploaded_pdfs_process_hf = st.file_uploader("Upload PDF files to process", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader_hf") + + if uploaded_pdfs_process_hf: + # Simplified: Process only the first page for demonstration + process_all_pages_pdf = st.checkbox("Process All Pages (can be slow/expensive)", value=False, key="pdf_process_all_hf") + pdf_prompt = st.text_area("Prompt for PDF Page Processing", "Extract the text content from this page.", key="pdf_process_prompt_hf") + + if st.button("Process Uploaded PDFs with HF", key="process_uploaded_pdfs_hf"): + if pdf_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: + st.error("Cannot process locally: No local model selected.") + else: + combined_text_output_hf = f"# HF PDF Processing Results ({'API' if pdf_use_api else 'Local'})\n\n" + total_pages_processed_hf = 0 + output_placeholder_hf = st.container() + + for pdf_file in uploaded_pdfs_process_hf: + output_placeholder_hf.markdown(f"--- \n### Processing: {pdf_file.name}") + # Read PDF bytes + pdf_bytes = pdf_file.read() + try: + doc = fitz.open("pdf", pdf_bytes) # Open from bytes + num_pages = len(doc) + pages_to_process = range(num_pages) if process_all_pages_pdf else range(min(1, num_pages)) # Limit to 1 unless checked + + output_placeholder_hf.info(f"Processing {len(pages_to_process)} of {num_pages} pages...") + doc_text = f"## File: {pdf_file.name}\n\n" + + for i in pages_to_process: + page_placeholder = output_placeholder_hf.empty() + page_placeholder.info(f"Processing Page {i + 1}/{num_pages}...") + page = doc[i] + pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) + img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) + + # Display image and process + cols_pdf = output_placeholder_hf.columns(2) + cols_pdf[0].image(img, caption=f"Page {i+1}", use_container_width=True) + with cols_pdf[1]: + # Use the new image processing function + # NOTE: This relies on the process_image_hf implementation + # which is currently basic/placeholder for local models. + with st.spinner("Processing page with HF model..."): + hf_text = process_image_hf(img, pdf_prompt, use_api=pdf_use_api) + st.text_area(f"Result (Page {i+1})", hf_text, height=250, key=f"pdf_hf_out_{pdf_file.name}_{i}") + + doc_text += f"### Page {i + 1}\n\n{hf_text}\n\n---\n\n" + total_pages_processed_hf += 1 + page_placeholder.empty() # Clear status message + + combined_text_output_hf += doc_text + doc.close() + + except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: + output_placeholder_hf.error(f"Error opening PDF {pdf_file.name}: {pdf_err}. Skipping.") + except Exception as e: + output_placeholder_hf.error(f"Error processing {pdf_file.name}: {str(e)}") + + if total_pages_processed_hf > 0: + st.markdown("--- \n### Combined Processing Results") + st.text_area("Full Output", combined_text_output_hf, height=400, key="combined_pdf_hf_output") + output_filename_pdf_hf = generate_filename("hf_processed_pdfs", "md") + try: + with open(output_filename_pdf_hf, "w", encoding="utf-8") as f: f.write(combined_text_output_hf) + st.success(f"Combined output saved to {output_filename_pdf_hf}") + st.markdown(get_download_link(output_filename_pdf_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True) + st.session_state['asset_checkboxes'][output_filename_pdf_hf] = False; update_gallery() + except IOError as e: st.error(f"Failed to save combined output file: {e}") + + +# --- Tab 6: Image Process (HF) --- +with tabs[5]: + st.header("Image Process with HF Models πŸ–ΌοΈ") + st.markdown("Upload images and process them using selected HF models (API or Local).") + + img_use_api = st.radio( + "Choose Processing Method", + ["Hugging Face API", "Loaded Local Model"], + key="img_process_source_hf", + horizontal=True + ) + if img_use_api == "Hugging Face API": + st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}") + else: + if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}") + else: st.warning("No local model selected.") + + img_prompt_hf = st.text_area("Prompt for Image Processing", "Describe this image in detail.", key="img_process_prompt_hf") + uploaded_images_process_hf = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader_hf") + + if uploaded_images_process_hf: + if st.button("Process Uploaded Images with HF", key="process_images_hf"): + if img_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: + st.error("Cannot process locally: No local model selected.") + else: + combined_img_text_hf = f"# HF Image Processing Results ({'API' if img_use_api else 'Local'})\n\n**Prompt:** {img_prompt_hf}\n\n---\n\n" + images_processed_hf = 0 + output_img_placeholder_hf = st.container() + + for img_file in uploaded_images_process_hf: + output_img_placeholder_hf.markdown(f"### Processing: {img_file.name}") + try: + img = Image.open(img_file) + cols_img_hf = output_img_placeholder_hf.columns(2) + cols_img_hf[0].image(img, caption=f"Input: {img_file.name}", use_container_width=True) + with cols_img_hf[1], st.spinner("Processing image with HF model..."): + # Use the new image processing function + hf_img_text = process_image_hf(img, img_prompt_hf, use_api=img_use_api) + st.text_area(f"Result", hf_img_text, height=300, key=f"img_hf_out_{img_file.name}") + + combined_img_text_hf += f"## Image: {img_file.name}\n\n{hf_img_text}\n\n---\n\n" + images_processed_hf += 1 + + except UnidentifiedImageError: output_img_placeholder_hf.error(f"Invalid Image: {img_file.name}. Skipping.") + except Exception as e: output_img_placeholder_hf.error(f"Error processing {img_file.name}: {str(e)}") + + if images_processed_hf > 0: + st.markdown("--- \n### Combined Processing Results") + st.text_area("Full Output", combined_img_text_hf, height=400, key="combined_img_hf_output") + output_filename_img_hf = generate_filename("hf_processed_images", "md") + try: + with open(output_filename_img_hf, "w", encoding="utf-8") as f: f.write(combined_img_text_hf) + st.success(f"Combined output saved to {output_filename_img_hf}") + st.markdown(get_download_link(output_filename_img_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True) + st.session_state['asset_checkboxes'][output_filename_img_hf] = False; update_gallery() + except IOError as e: st.error(f"Failed to save combined output file: {e}") + + +# --- Tab 7: Text Process (HF) --- +with tabs[6]: + st.header("Text Process with HF Models πŸ“") + st.markdown("Process Markdown (.md) or Text (.txt) files using selected HF models (API or Local).") + + text_use_api = st.radio( + "Choose Processing Method", + ["Hugging Face API", "Loaded Local Model"], + key="text_process_source_hf", + horizontal=True + ) + if text_use_api == "Hugging Face API": + st.info(f"Using API Model: {st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model}") + else: + if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}") + else: st.warning("No local model selected.") + + text_files_hf = get_gallery_files(['md', 'txt']) + if not text_files_hf: + st.warning("No .md or .txt files in gallery to process.") + else: + selected_text_file_hf = st.selectbox( + "Select Text/MD File to Process", + options=[""] + text_files_hf, + format_func=lambda x: os.path.basename(x) if x else "Select a file...", + key="text_process_select_hf" + ) + + if selected_text_file_hf: + st.write(f"Selected: {os.path.basename(selected_text_file_hf)}") + try: + with open(selected_text_file_hf, "r", encoding="utf-8", errors='ignore') as f: + content_text_hf = f.read() + st.text_area("File Content Preview", content_text_hf[:1000] + ("..." if len(content_text_hf) > 1000 else ""), height=200, key="text_content_preview_hf") + + prompt_text_hf = st.text_area( + "Enter Prompt for this File", + "Summarize the key points of this text.", + key="text_individual_prompt_hf" + ) + + if st.button(f"Process '{os.path.basename(selected_text_file_hf)}' with HF", key=f"process_text_hf_btn"): + if text_use_api == "Loaded Local Model" and not st.session_state.selected_local_model_path: + st.error("Cannot process locally: No local model selected.") + else: + with st.spinner("Processing text with HF model..."): + result_text_processed = process_text_hf(content_text_hf, prompt_text_hf, use_api=text_use_api) + + st.markdown("### Processing Result") + st.markdown(result_text_processed) # Display result + + output_filename_text_hf = generate_filename(f"hf_processed_{os.path.splitext(os.path.basename(selected_text_file_hf))[0]}", "md") + try: + with open(output_filename_text_hf, "w", encoding="utf-8") as f: f.write(result_text_processed) + st.success(f"Result saved to {output_filename_text_hf}") + st.markdown(get_download_link(output_filename_text_hf, "text/markdown", "Download Result MD"), unsafe_allow_html=True) + st.session_state['asset_checkboxes'][output_filename_text_hf] = False; update_gallery() + except IOError as e: st.error(f"Failed to save result file: {e}") + + except FileNotFoundError: st.error("Selected file not found.") + except Exception as e: st.error(f"Error reading file: {e}") + + +# --- Tab 8: Test OCR (HF) --- +with tabs[7]: + st.header("Test OCR with HF Models πŸ”") + st.markdown("Select an image/PDF and run OCR using HF models (API or Local - requires suitable local model).") + + ocr_use_api = st.radio( + "Choose OCR Method", + ["Hugging Face API (Basic Captioning/OCR)", "Loaded Local OCR Model"], + key="ocr_source_hf", + horizontal=True, + help="API uses basic image-to-text. Local requires a dedicated OCR model (e.g., TrOCR) to be loaded." + ) + if ocr_use_api == "Loaded Local OCR Model": + if st.session_state.selected_local_model_path: + model_type = st.session_state.local_models.get(st.session_state.selected_local_model_path,{}).get('type') + if model_type != 'ocr': + st.warning(f"Selected local model ({os.path.basename(st.session_state.selected_local_model_path)}) is type '{model_type}', not 'ocr'. Results may be poor.") + else: + st.info(f"Using Local OCR Model: {os.path.basename(st.session_state.selected_local_model_path)}") + else: st.warning("No local model selected.") + + gallery_files_ocr_hf = get_gallery_files(['png', 'jpg', 'jpeg', 'pdf']) + if not gallery_files_ocr_hf: + st.warning("No images or PDFs in gallery.") + else: + selected_file_ocr_hf = st.selectbox( + "Select Image or PDF from Gallery for OCR", + options=[""] + gallery_files_ocr_hf, + format_func=lambda x: os.path.basename(x) if x else "Select a file...", + key="ocr_select_file_hf" + ) + + if selected_file_ocr_hf: + st.write(f"Selected: {os.path.basename(selected_file_ocr_hf)}") + file_ext_ocr_hf = os.path.splitext(selected_file_ocr_hf)[1].lower() + image_to_ocr_hf = None; page_info_hf = "" + + try: + if file_ext_ocr_hf in ['.png', '.jpg', '.jpeg']: image_to_ocr_hf = Image.open(selected_file_ocr_hf) + elif file_ext_ocr_hf == '.pdf': + doc = fitz.open(selected_file_ocr_hf) + if len(doc) > 0: pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image_to_ocr_hf = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); page_info_hf = " (Page 1)" + else: st.warning("Selected PDF is empty.") + doc.close() + + if image_to_ocr_hf: + st.image(image_to_ocr_hf, caption=f"Image for OCR{page_info_hf}", use_container_width=True) + if st.button("Run HF OCR on this Image πŸš€", key="ocr_run_button_hf"): + if ocr_use_api == "Loaded Local OCR Model" and not st.session_state.selected_local_model_path: + st.error("Cannot run locally: No local model selected.") + else: + output_ocr_file_hf = generate_filename(f"hf_ocr_{os.path.splitext(os.path.basename(selected_file_ocr_hf))[0]}", "txt") + st.session_state['processing']['ocr'] = True + with st.spinner("Performing OCR with HF model..."): + ocr_result_hf = asyncio.run(process_hf_ocr(image_to_ocr_hf, output_ocr_file_hf, use_api=ocr_use_api)) + st.session_state['processing']['ocr'] = False + + st.text_area("OCR Result", ocr_result_hf, height=300, key="ocr_result_display_hf") + if ocr_result_hf and not ocr_result_hf.startswith("Error") and not ocr_result_hf.startswith("["): + entry = f"HF OCR: {selected_file_ocr_hf}{page_info_hf} -> {output_ocr_file_hf}" + st.session_state['history'].append(entry) + if len(ocr_result_hf) > 5: # Minimal check + st.success(f"OCR output saved to {output_ocr_file_hf}") + st.markdown(get_download_link(output_ocr_file_hf, "text/plain", "Download OCR Text"), unsafe_allow_html=True) + st.session_state['asset_checkboxes'][output_ocr_file_hf] = False; update_gallery() + else: st.warning("OCR output seems short/empty.") + else: st.error(f"OCR failed. {ocr_result_hf}") + + except Exception as e: st.error(f"Error loading file for OCR: {e}") + +# --- Tab 9: Test Image Gen (Diffusers) --- +with tabs[8]: + st.header("Test Image Generation (Diffusers) 🎨") + st.markdown("Generate images using Stable Diffusion models loaded locally via the Diffusers library.") + + if not _diffusers_available: + st.error("Diffusers library is required for image generation.") + else: + # Select from locally downloaded *diffusion* models + local_diffusion_paths = get_local_model_paths("diffusion") + if not local_diffusion_paths: + st.warning("No local diffusion models found. Download one using the 'Build Titan' tab.") + selected_diffusion_model_path = None + else: + selected_diffusion_model_path = st.selectbox( + "Select Local Diffusion Model", + options=[""] + local_diffusion_paths, + format_func=lambda x: os.path.basename(x) if x else "Select...", + key="imggen_diffusion_model_select" + ) + + prompt_imggen_diff = st.text_area("Image Generation Prompt", "A photorealistic cat wearing sunglasses, studio lighting", key="imggen_prompt_diff") + neg_prompt_imggen_diff = st.text_area("Negative Prompt (Optional)", "ugly, deformed, blurry, low quality", key="imggen_neg_prompt_diff") + steps_imggen_diff = st.slider("Inference Steps", 10, 100, 25, key="imggen_steps") + guidance_imggen_diff = st.slider("Guidance Scale", 1.0, 20.0, 7.5, step=0.5, key="imggen_guidance") + + if st.button("Generate Image πŸš€", key="imggen_run_button_diff", disabled=not selected_diffusion_model_path): + if not prompt_imggen_diff: st.warning("Please enter a prompt.") + else: + status_imggen = st.empty() + try: + # Load pipeline from saved path on demand + status_imggen.info(f"Loading diffusion pipeline: {os.path.basename(selected_diffusion_model_path)}...") + # Determine device + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float16 if torch.cuda.is_available() else torch.float32 # Use float16 on GPU if available + pipe = StableDiffusionPipeline.from_pretrained(selected_diffusion_model_path, torch_dtype=dtype).to(device) + pipe.safety_checker = None # Optional: Disable safety checker if needed + + status_imggen.info(f"Generating image on {device} ({dtype})...") + start_gen_time = time.time() + + # Generate using the pipeline + gen_output = pipe( + prompt=prompt_imggen_diff, + negative_prompt=neg_prompt_imggen_diff if neg_prompt_imggen_diff else None, + num_inference_steps=steps_imggen_diff, + guidance_scale=guidance_imggen_diff, + # Add seed if desired: generator=torch.Generator(device=device).manual_seed(your_seed) + ) + gen_image = gen_output.images[0] + + elapsed_gen = int(time.time() - start_gen_time) + status_imggen.success(f"Image generated in {elapsed_gen}s!") + + # Save and display + output_imggen_file_diff = generate_filename("diffusion_gen", "png") + gen_image.save(output_imggen_file_diff) + st.image(gen_image, caption=f"Generated: {output_imggen_file_diff}", use_container_width=True) + st.markdown(get_download_link(output_imggen_file_diff, "image/png", "Download Generated Image"), unsafe_allow_html=True) + st.session_state['asset_checkboxes'][output_imggen_file_diff] = False; update_gallery() + st.session_state['history'].append(f"Diffusion Gen: '{prompt_imggen_diff[:30]}...' -> {output_imggen_file_diff}") + + except ImportError: st.error("Diffusers or Torch library not found.") + except Exception as e: + st.error(f"Image generation failed: {e}") + logger.error(f"Diffusion generation failed for {selected_diffusion_model_path}: {e}") + finally: + # Clear pipeline from memory? (Optional, depends on memory usage) + if 'pipe' in locals(): del pipe; torch.cuda.empty_cache() if torch.cuda.is_available() else None + + +# --- Tab 10: Character Editor (Keep from previous merge) --- +with tabs[9]: + # ... (Code from previous merge for this tab) ... + st.header("Character Editor πŸ§‘β€πŸŽ¨") + st.subheader("Create Your Character") + load_characters(); existing_char_names = [c['name'] for c in st.session_state.get('characters', [])] + form_key = f"character_form_{st.session_state.get('char_form_reset_key', 0)}" + with st.form(key=form_key): + st.markdown("**Create New Character**") + if st.form_submit_button("Randomize Content 🎲"): st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1; st.rerun() + rand_name, rand_gender, rand_intro, rand_greeting = randomize_character_content() + name_char = st.text_input("Name (3-25 chars...)", value=rand_name, max_chars=25, key="char_name_input") + gender_char = st.radio("Gender", ["Male", "Female"], index=["Male", "Female"].index(rand_gender), key="char_gender_radio") + intro_char = st.text_area("Intro (Public description)", value=rand_intro, max_chars=300, height=100, key="char_intro_area") + greeting_char = st.text_area("Greeting (First message)", value=rand_greeting, max_chars=300, height=100, key="char_greeting_area") + tags_char = st.text_input("Tags (comma-separated)", "OC, friendly", key="char_tags_input") + submitted = st.form_submit_button("Create Character ✨") + if submitted: + error = False + if not (3 <= len(name_char) <= 25): st.error("Name must be 3-25 characters."); error = True + if not re.match(r'^[a-zA-Z0-9 _-]+$', name_char): st.error("Name contains invalid characters."); error = True + if name_char in existing_char_names: st.error(f"Name '{name_char}' already exists!"); error = True + if not intro_char or not greeting_char: st.error("Intro/Greeting cannot be empty."); error = True + if not error: + tag_list = [tag.strip() for tag in tags_char.split(',') if tag.strip()] + character_data = {"name": name_char, "gender": gender_char, "intro": intro_char, "greeting": greeting_char, "created_at": datetime.now(pytz.timezone("US/Central")).strftime('%Y-%m-%d %H:%M:%S %Z'), "tags": tag_list} + if save_character(character_data): + st.success(f"Character '{name_char}' created!"); st.session_state['char_form_reset_key'] = st.session_state.get('char_form_reset_key', 0) + 1; st.rerun() + +# --- Tab 11: Character Gallery (Keep from previous merge) --- +with tabs[10]: + # ... (Code from previous merge for this tab) ... + st.header("Character Gallery πŸ–ΌοΈ") + load_characters(); characters_list = st.session_state.get('characters', []) + if not characters_list: st.warning("No characters created yet.") + else: + st.subheader(f"Your Characters ({len(characters_list)})") + search_term = st.text_input("Search Characters by Name", key="char_gallery_search") + if search_term: characters_list = [c for c in characters_list if search_term.lower() in c['name'].lower()] + cols_char_gallery = st.columns(3); chars_to_delete = [] + for idx, char in enumerate(characters_list): + with cols_char_gallery[idx % 3], st.container(border=True): + st.markdown(f"**{char['name']}**"); st.caption(f"Gender: {char.get('gender', 'N/A')}") + st.markdown("**Intro:**"); st.markdown(f"> {char.get('intro', '')}") + st.markdown("**Greeting:**"); st.markdown(f"> {char.get('greeting', '')}") + st.caption(f"Tags: {', '.join(char.get('tags', ['N/A']))}"); st.caption(f"Created: {char.get('created_at', 'N/A')}") + delete_key_char = f"delete_char_{char['name']}_{idx}"; + if st.button(f"Delete {char['name']}", key=delete_key_char, type="primary"): chars_to_delete.append(char['name']) + if chars_to_delete: + current_characters = st.session_state.get('characters', []); updated_characters = [c for c in current_characters if c['name'] not in chars_to_delete] + st.session_state['characters'] = updated_characters + try: + with open("characters.json", "w", encoding='utf-8') as f: json.dump(updated_characters, f, indent=2) + logger.info(f"Deleted characters: {', '.join(chars_to_delete)}"); st.success(f"Deleted characters: {', '.join(chars_to_delete)}"); st.rerun() + except IOError as e: logger.error(f"Failed to save characters.json after deletion: {e}"); st.error("Failed to update character file.") + +# --- Footer and Persistent Sidebar Elements ------------ + +# Update Sidebar Gallery (Call this at the end to reflect all changes) +update_gallery() + +# Action Logs in Sidebar +st.sidebar.subheader("Action Logs πŸ“œ") +log_expander = st.sidebar.expander("View Logs", expanded=False) +with log_expander: + log_text = "\n".join([f"{record.asctime} - {record.levelname} - {record.message}" for record in log_records[-20:]]) + st.code(log_text, language='log') + +# History in Sidebar +st.sidebar.subheader("Session History πŸ“œ") +history_expander = st.sidebar.expander("View History", expanded=False) +with history_expander: + for entry in reversed(st.session_state.get("history", [])): + if entry: history_expander.write(f"- {entry}") + +st.sidebar.markdown("---") +st.sidebar.info("Using Hugging Face models for AI tasks.") +st.sidebar.caption("App Modified by AI Assistant") \ No newline at end of file