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