awacke1 commited on
Commit
31a2b78
·
verified ·
1 Parent(s): 1a09829

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +105 -1369
app.py CHANGED
@@ -1,1379 +1,115 @@
1
- # --- Combined Imports ------------------------------------
2
- import io
3
  import os
4
- import re
5
- import base64
6
  import glob
7
- import logging
8
- import random
9
- import shutil
10
- import time
11
- import zipfile
12
- import json
13
- import asyncio
14
- import aiofiles
15
-
16
- from datetime import datetime
17
- from collections import Counter
18
- from dataclasses import dataclass, field
19
- from io import BytesIO
20
- from typing import Optional, List, Dict, Any
21
-
22
  import pandas as pd
23
- import pytz
24
- import streamlit as st
25
- from PIL import Image, ImageDraw, UnidentifiedImageError # Added ImageDraw and UnidentifiedImageError
 
26
  from reportlab.pdfgen import canvas
27
  from reportlab.lib.utils import ImageReader
28
- from reportlab.lib.pagesizes import letter # Default page size
29
- from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PlatypusImage, PageBreak, Preformatted
30
- from reportlab.lib.styles import getSampleStyleSheet
31
- from reportlab.lib.units import inch
32
- from reportlab.lib.enums import TA_CENTER, TA_LEFT # For text alignment
33
- import fitz # PyMuPDF
34
-
35
- # --- Hugging Face Imports ---
36
- from huggingface_hub import InferenceClient, HfApi, list_models
37
- from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Import specific exceptions
38
 
39
- # --- App Configuration -----------------------------------
40
  st.set_page_config(
41
- page_title="Vision & Layout Titans (HF) 🚀🖼️",
42
  page_icon="🤖",
43
- layout="wide",
44
- initial_sidebar_state="expanded",
45
- menu_items={
46
- 'Get Help': 'https://huggingface.co/docs',
47
- 'Report a Bug': None, # Replace with your bug report link if desired
48
- 'About': "Combined App: PDF Layout Generator + Hugging Face Powered AI Tools 🌌"
49
- }
50
  )
51
 
52
-
53
- # Conditional imports for optional/heavy libraries
54
- try:
55
- import torch
56
- from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForVision2Seq, AutoModelForImageToWaveform, pipeline
57
- # Add more AutoModel classes as needed for different tasks (Vision, OCR, etc.)
58
- _transformers_available = True
59
- except ImportError:
60
- _transformers_available = False
61
- # Place warning inside main app area if sidebar isn't ready
62
- # st.sidebar.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.")
63
-
64
- try:
65
- from diffusers import StableDiffusionPipeline
66
- _diffusers_available = True
67
- except ImportError:
68
- _diffusers_available = False
69
- # Don't show warning if transformers also missing, handled above
70
- # if _transformers_available:
71
- # st.sidebar.warning("Diffusers library not found. Diffusion model features disabled.")
72
-
73
-
74
- import requests # Keep requests import
75
-
76
- # --- Logging Setup ---------------------------------------
77
- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
78
- logger = logging.getLogger(__name__)
79
- log_records = []
80
- class LogCaptureHandler(logging.Handler):
81
- def emit(self, record):
82
- # Limit stored logs to avoid memory issues
83
- if len(log_records) > 200:
84
- log_records.pop(0)
85
- log_records.append(record)
86
- logger.addHandler(LogCaptureHandler())
87
-
88
- # --- Environment Variables & Constants -------------------
89
- HF_TOKEN = os.getenv("HF_TOKEN")
90
- DEFAULT_PROVIDER = "hf-inference"
91
- # Model List (curated, similar to Gradio example) - can be updated
92
- FEATURED_MODELS_LIST = [
93
- "meta-llama/Meta-Llama-3.1-8B-Instruct", # Updated Llama model
94
- "mistralai/Mistral-7B-Instruct-v0.3",
95
- "google/gemma-2-9b-it", # Added Gemma 2
96
- "Qwen/Qwen2-7B-Instruct", # Added Qwen2
97
- "microsoft/Phi-3-mini-4k-instruct",
98
- "HuggingFaceH4/zephyr-7b-beta",
99
- "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", # Larger Mixture of Experts
100
- # Add a smaller option
101
- "HuggingFaceTB/SmolLM-1.7B-Instruct"
102
- ]
103
- # Add common vision models if planning local loading
104
- VISION_MODELS_LIST = [
105
- "Salesforce/blip-image-captioning-large",
106
- "microsoft/trocr-large-handwritten", # OCR model
107
- "llava-hf/llava-1.5-7b-hf", # Vision Language Model
108
- "google/vit-base-patch16-224", # Basic Vision Transformer
109
- ]
110
- DIFFUSION_MODELS_LIST = [
111
- "stabilityai/stable-diffusion-xl-base-1.0", # Common SDXL
112
- "runwayml/stable-diffusion-v1-5", # Classic SD 1.5
113
- "OFA-Sys/small-stable-diffusion-v0", # Tiny diffusion
114
- ]
115
-
116
-
117
- # --- Session State Initialization (Combined & Updated) ---
118
- # Combined PDF Generator specific (replaces layout specific)
119
- st.session_state.setdefault('combined_pdf_sources', []) # List of dicts {'filepath': path, 'type': type}
120
-
121
- # General App State
122
- st.session_state.setdefault('history', [])
123
- st.session_state.setdefault('processing', {})
124
- st.session_state.setdefault('asset_checkboxes', {})
125
- st.session_state.setdefault('downloaded_pdfs', {})
126
- st.session_state.setdefault('unique_counter', 0)
127
- st.session_state.setdefault('cam0_file', None)
128
- st.session_state.setdefault('cam1_file', None)
129
- st.session_state.setdefault('characters', [])
130
- st.session_state.setdefault('char_form_reset_key', 0) # For character form reset
131
- # Removed gallery_size state - no longer used
132
- # st.session_state.setdefault('gallery_size', 10)
133
-
134
- # --- Hugging Face & Local Model State ---
135
- st.session_state.setdefault('hf_inference_client', None) # Store initialized client
136
- st.session_state.setdefault('hf_provider', DEFAULT_PROVIDER)
137
- st.session_state.setdefault('hf_custom_key', "")
138
- st.session_state.setdefault('hf_selected_api_model', FEATURED_MODELS_LIST[0]) # Default API model
139
- st.session_state.setdefault('hf_custom_api_model', "") # User override for API model
140
-
141
- # Local Model Management
142
- st.session_state.setdefault('local_models', {}) # Dict to store loaded models: {'path': {'model': obj, 'tokenizer/proc': obj, 'type': 'causal/vision/etc'}}
143
- st.session_state.setdefault('selected_local_model_path', None) # Path of the currently active local model
144
-
145
- # Inference Parameters (shared for API and local where applicable)
146
- st.session_state.setdefault('gen_max_tokens', 512)
147
- st.session_state.setdefault('gen_temperature', 0.7)
148
- st.session_state.setdefault('gen_top_p', 0.95)
149
- st.session_state.setdefault('gen_frequency_penalty', 0.0) # Corresponds to repetition_penalty=1.0
150
- st.session_state.setdefault('gen_seed', -1) # -1 for random
151
-
152
- # Removed asset_gallery_container - render directly in sidebar
153
- # if 'asset_gallery_container' not in st.session_state:
154
- # st.session_state['asset_gallery_container'] = st.sidebar.empty()
155
-
156
- # --- Dataclasses (Refined for Local Models) -------------
157
- @dataclass
158
- class LocalModelConfig:
159
- name: str # User-defined local name
160
- hf_id: str # Hugging Face model ID used for download
161
- model_type: str # 'causal', 'vision', 'diffusion', 'ocr', etc.
162
- size_category: str = "unknown" # e.g., 'small', 'medium', 'large'
163
- domain: Optional[str] = None
164
- local_path: str = field(init=False) # Path where it's saved
165
-
166
- def __post_init__(self):
167
- # Define local path based on type and name
168
- type_folder = f"{self.model_type}_models"
169
- safe_name = re.sub(r'[^\w\-]+', '_', self.name) # Sanitize name for path
170
- self.local_path = os.path.join(type_folder, safe_name)
171
-
172
- def get_full_path(self):
173
- return os.path.abspath(self.local_path)
174
-
175
- # (Keep DiffusionConfig if still using diffusers library separately)
176
- @dataclass
177
- class DiffusionConfig: # Kept for clarity in diffusion tab if needed
178
- name: str
179
- base_model: str
180
- size: str
181
- domain: Optional[str] = None
182
- @property
183
- def model_path(self):
184
- # Ensure diffusion models are saved in their own distinct top-level folder
185
- return f"diffusion_models/{re.sub(r'[^w-]+', '_', self.name)}"
186
-
187
-
188
- # --- Helper Functions (Combined and refined) -------------
189
- def generate_filename(sequence, ext="png"):
190
- timestamp = time.strftime('%Y%m%d_%H%M%S')
191
- safe_sequence = re.sub(r'[^\w\-]+', '_', str(sequence))
192
- return f"{safe_sequence}_{timestamp}.{ext}"
193
-
194
- def pdf_url_to_filename(url):
195
- name = re.sub(r'^https?://', '', url)
196
- name = re.sub(r'[<>:"/\\|?*]', '_', name)
197
- return name[:100] + ".pdf" # Limit length
198
-
199
- def get_download_link(file_path, mime_type="application/octet-stream", label="Download"):
200
- if not os.path.exists(file_path): return f"{label} (File not found)"
201
- try:
202
- with open(file_path, "rb") as f: file_bytes = f.read()
203
- b64 = base64.b64encode(file_bytes).decode()
204
- return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
205
- except Exception as e:
206
- logger.error(f"Error creating download link for {file_path}: {e}")
207
- return f"{label} (Error)"
208
-
209
- def zip_directory(directory_path, zip_path):
210
- with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
211
- for root, _, files in os.walk(directory_path):
212
- for file in files:
213
- file_path = os.path.join(root, file)
214
- zipf.write(file_path, os.path.relpath(file_path, os.path.dirname(directory_path)))
215
-
216
- def get_local_model_paths(model_type="causal"):
217
- """Gets paths of locally saved models of a specific type."""
218
- pattern = f"{model_type}_models/*"
219
- dirs = [d for d in glob.glob(pattern) if os.path.isdir(d)]
220
- return dirs
221
-
222
- def get_gallery_files(file_types=("png", "pdf", "jpg", "jpeg", "md", "txt")):
223
- """Gets all files with specified extensions in the current directory."""
224
- all_files = set()
225
- for ext in file_types:
226
- # Ensure the glob pattern correctly targets files in the script's directory
227
- all_files.update(glob.glob(f"./*.{ext.lower()}")) # Use ./* for current dir
228
- all_files.update(glob.glob(f"./*.{ext.upper()}"))
229
- # Convert to list and remove potential './' prefix for cleaner display
230
- return sorted([os.path.normpath(f) for f in all_files])
231
-
232
- def get_pdf_files():
233
- # Use get_gallery_files to find PDFs
234
- return get_gallery_files(['pdf'])
235
-
236
- def download_pdf(url, output_path):
237
- try:
238
- headers = {'User-Agent': 'Mozilla/5.0'}
239
- response = requests.get(url, stream=True, timeout=20, headers=headers)
240
- response.raise_for_status()
241
- with open(output_path, "wb") as f:
242
- for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
243
- logger.info(f"Successfully downloaded {url} to {output_path}")
244
- return True
245
- except requests.exceptions.RequestException as e:
246
- logger.error(f"Failed to download {url}: {e}")
247
- if os.path.exists(output_path):
248
- try:
249
- os.remove(output_path)
250
- logger.info(f"Removed partially downloaded file: {output_path}")
251
- except OSError as remove_error:
252
- logger.error(f"Error removing partial file {output_path}: {remove_error}")
253
- except Exception as general_remove_error:
254
- logger.error(f"General error removing partial file {output_path}: {general_remove_error}")
255
- return False
256
- except Exception as e:
257
- logger.error(f"An unexpected error occurred during download of {url}: {e}")
258
- if os.path.exists(output_path):
259
- try:
260
- os.remove(output_path)
261
- logger.info(f"Removed file after unexpected error: {output_path}")
262
- except OSError as remove_error:
263
- logger.error(f"Error removing file after unexpected error {output_path}: {remove_error}")
264
- except Exception as general_remove_error:
265
- logger.error(f"General error removing file after unexpected error {output_path}: {general_remove_error}")
266
- return False
267
-
268
- async def process_pdf_snapshot(pdf_path, mode="single", resolution_factor=2.0):
269
- start_time = time.time()
270
- # Use a placeholder within the main app area for status during async operations
271
- status_placeholder = st.empty()
272
- status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... (0s)")
273
- output_files = []
274
- try:
275
- doc = fitz.open(pdf_path)
276
- matrix = fitz.Matrix(resolution_factor, resolution_factor)
277
- num_pages_to_process = 0
278
- if mode == "single": num_pages_to_process = min(1, len(doc))
279
- elif mode == "twopage": num_pages_to_process = min(2, len(doc))
280
- elif mode == "allpages": num_pages_to_process = len(doc)
281
-
282
- for i in range(num_pages_to_process):
283
- page_start_time = time.time()
284
- page = doc.load_page(i) # Use load_page for efficiency
285
- pix = page.get_pixmap(matrix=matrix)
286
- base_name = os.path.splitext(os.path.basename(pdf_path))[0]
287
- output_file = generate_filename(f"{base_name}_pg{i+1}_{mode}", "png")
288
-
289
- # Ensure output path is valid before saving
290
- output_dir = os.path.dirname(output_file) or "."
291
- if not os.path.exists(output_dir): os.makedirs(output_dir)
292
-
293
- await asyncio.to_thread(pix.save, output_file)
294
- output_files.append(output_file)
295
- elapsed_page = int(time.time() - page_start_time)
296
- status_placeholder.text(f"Processing PDF Snapshot ({mode}, Res: {resolution_factor}x)... Page {i+1}/{num_pages_to_process} done ({elapsed_page}s)")
297
- await asyncio.sleep(0.01)
298
-
299
- doc.close()
300
- elapsed = int(time.time() - start_time)
301
- status_placeholder.success(f"PDF Snapshot ({mode}, {len(output_files)} files) completed in {elapsed}s!")
302
- return output_files
303
- except Exception as e:
304
- logger.error(f"Failed to process PDF snapshot for {pdf_path}: {e}", exc_info=True) # Add traceback
305
- status_placeholder.error(f"Failed to process PDF {os.path.basename(pdf_path)}: {e}")
306
- # Clean up any files created before the error
307
- for f in output_files:
308
- if os.path.exists(f):
309
- try: os.remove(f)
310
- except: pass
311
- return []
312
-
313
-
314
- # --- HF Inference Client Management ---
315
- def get_hf_client() -> Optional[InferenceClient]:
316
- """Gets or initializes the Hugging Face Inference Client based on session state."""
317
- provider = st.session_state.hf_provider
318
- custom_key = st.session_state.hf_custom_key.strip()
319
- token_to_use = custom_key if custom_key else HF_TOKEN
320
-
321
- if not token_to_use and provider != "hf-inference":
322
- # Don't show error here, let caller handle it if client is needed
323
- # st.error(f"Provider '{provider}' requires a Hugging Face API token...")
324
- return None
325
- if provider == "hf-inference" and not token_to_use:
326
- logger.warning("Using hf-inference provider without a token. Rate limits may apply.")
327
- token_to_use = None # Explicitly set to None for public inference API
328
-
329
- # Check if client needs re-initialization
330
- current_client = st.session_state.get('hf_inference_client')
331
- needs_reinit = True
332
- if current_client:
333
- # Compare provider and token status more carefully
334
- current_token = getattr(current_client, '_token', None) # Access internal token if exists
335
- current_provider = getattr(current_client, 'provider', None) # Access provider if exists
336
-
337
- token_matches = (token_to_use == current_token)
338
- provider_matches = (provider == current_provider)
339
-
340
- if token_matches and provider_matches:
341
- needs_reinit = False
342
-
343
- if needs_reinit:
344
- try:
345
- logger.info(f"Initializing InferenceClient for provider: {provider}. Token source: {'Custom Key' if custom_key else ('HF_TOKEN' if HF_TOKEN else 'None')}")
346
- st.session_state.hf_inference_client = InferenceClient(model=None, token=token_to_use, provider=provider) # Init without model initially
347
- # Store provider on client instance if possible (check InferenceClient structure or assume it's handled internally)
348
- setattr(st.session_state.hf_inference_client, 'provider', provider) # Explicitly store provider for re-init check
349
- setattr(st.session_state.hf_inference_client, '_token', token_to_use) # Explicitly store token for re-init check
350
- logger.info("InferenceClient initialized successfully.")
351
- except Exception as e:
352
- st.error(f"Failed to initialize Hugging Face client for provider {provider}: {e}")
353
- logger.error(f"InferenceClient initialization failed: {e}")
354
- st.session_state.hf_inference_client = None
355
-
356
- return st.session_state.hf_inference_client
357
-
358
- # --- HF/Local Model Processing Functions ---
359
- def process_text_hf(text: str, prompt: str, use_api: bool) -> str:
360
- """Processes text using either HF Inference API or a loaded local model."""
361
- status_placeholder = st.empty()
362
- start_time = time.time()
363
- result_text = ""
364
-
365
- params = {
366
- "max_new_tokens": st.session_state.gen_max_tokens,
367
- "temperature": st.session_state.gen_temperature,
368
- "top_p": st.session_state.gen_top_p,
369
- "repetition_penalty": st.session_state.gen_frequency_penalty, # Keep user value, adjust name below if needed
370
- }
371
- seed = st.session_state.gen_seed
372
- if seed != -1: params["seed"] = seed
373
-
374
- system_prompt = "You are a helpful assistant. Process the following text based on the user's request."
375
- full_prompt = f"{prompt}\n\n---\n\n{text}"
376
- messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": full_prompt}]
377
-
378
- if use_api:
379
- status_placeholder.info("Processing text using Hugging Face API...")
380
- client = get_hf_client()
381
- if not client: return "Error: Hugging Face client not configured/available."
382
- model_id = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
383
- if not model_id: return "Error: No Hugging Face API model specified."
384
- status_placeholder.info(f"Using API Model: {model_id}")
385
- try:
386
- # Ensure repetition_penalty is passed correctly if supported
387
- api_params = {
388
- "max_tokens": params['max_new_tokens'],
389
- "temperature": params['temperature'],
390
- "top_p": params['top_p'],
391
- "repetition_penalty": params.get('repetition_penalty') # Check if API uses this name
392
- }
393
- if 'seed' in params: api_params['seed'] = params['seed']
394
-
395
- response = client.chat_completion(model=model_id, messages=messages, **api_params)
396
- result_text = response.choices[0].message.content or ""
397
- logger.info(f"HF API text processing successful for model {model_id}.")
398
- except Exception as e:
399
- logger.error(f"HF API text processing failed for model {model_id}: {e}", exc_info=True)
400
- result_text = f"Error during Hugging Face API inference: {str(e)}"
401
- else:
402
- status_placeholder.info("Processing text using local model...")
403
- if not _transformers_available: return "Error: Transformers library not available."
404
- model_path = st.session_state.get('selected_local_model_path')
405
- if not model_path or model_path not in st.session_state.get('local_models', {}): return "Error: No suitable local model selected/loaded."
406
- local_model_data = st.session_state['local_models'][model_path]
407
- if local_model_data.get('type') != 'causal': return f"Error: Loaded model '{os.path.basename(model_path)}' is not a Causal LM."
408
- status_placeholder.info(f"Using Local Model: {os.path.basename(model_path)}")
409
- model = local_model_data.get('model')
410
- tokenizer = local_model_data.get('tokenizer')
411
- if not model or not tokenizer: return f"Error: Model/tokenizer not found for {os.path.basename(model_path)}."
412
- try:
413
- try: prompt_for_model = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
414
- except: logger.warning(f"Chat template failed for {model_path}. Using basic format."); prompt_for_model = f"System: {system_prompt}\nUser: {full_prompt}\nAssistant:"
415
- inputs = tokenizer(prompt_for_model, return_tensors="pt", padding=True, truncation=True, max_length=2048).to(model.device) # Increased context slightly
416
- generate_params = {
417
- "max_new_tokens": params['max_new_tokens'],
418
- "temperature": params['temperature'],
419
- "top_p": params['top_p'],
420
- "repetition_penalty": params.get('repetition_penalty', 1.0),
421
- "do_sample": True if params['temperature'] > 0.01 else False, # Sample if temp > 0.01
422
- "pad_token_id": tokenizer.eos_token_id
423
- }
424
- with torch.no_grad(): outputs = model.generate(**inputs, **generate_params)
425
- input_length = inputs['input_ids'].shape[1]; generated_ids = outputs[0][input_length:]
426
- result_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
427
- logger.info(f"Local text processing successful for model {model_path}.")
428
- except Exception as e:
429
- logger.error(f"Local text processing failed for model {model_path}: {e}", exc_info=True)
430
- result_text = f"Error during local model inference: {str(e)}"
431
-
432
- elapsed = int(time.time() - start_time)
433
- status_placeholder.success(f"Text processing completed in {elapsed}s.")
434
- return result_text
435
-
436
- def process_image_hf(image: Image.Image, prompt: str, use_api: bool) -> str:
437
- """Processes an image using either HF Inference API or a local model."""
438
- status_placeholder = st.empty()
439
- start_time = time.time()
440
- result_text = "[Image processing requires specific Vision model implementation]"
441
-
442
- if use_api:
443
- status_placeholder.info("Processing image using Hugging Face API (Image-to-Text)...")
444
- client = get_hf_client()
445
- if not client: return "Error: HF client not configured."
446
- buffered = BytesIO(); image.save(buffered, format="PNG"); img_bytes = buffered.getvalue()
447
- try:
448
- captioning_model_id = "Salesforce/blip-image-captioning-large" # Default captioner
449
- vqa_model_id = "llava-hf/llava-1.5-7b-hf" # Default VQA - MAY REQUIRE DIFFERENT CLIENT CALL
450
- # Decide whether to use captioning or VQA based on prompt? Simple approach: captioning.
451
- status_placeholder.info(f"Using API Image-to-Text Model: {captioning_model_id}")
452
- response_list = client.image_to_text(data=img_bytes, model=captioning_model_id)
453
- if response_list and 'generated_text' in response_list[0]:
454
- result_text = f"API Caption: {response_list[0]['generated_text']}\n(Prompt '{prompt}' likely ignored by this API endpoint)"
455
- logger.info(f"HF API image captioning successful for model {captioning_model_id}.")
456
- else: result_text = "Error: Unexpected response format from image-to-text API."; logger.warning(f"Unexpected API response: {response_list}")
457
- except Exception as e: logger.error(f"HF API image processing failed: {e}"); result_text = f"Error during HF API image inference: {str(e)}"
458
- else:
459
- status_placeholder.info("Processing image using local model...")
460
- if not _transformers_available: return "Error: Transformers library needed."
461
- model_path = st.session_state.get('selected_local_model_path')
462
- if not model_path or model_path not in st.session_state.get('local_models', {}): return "Error: No suitable local model selected/loaded."
463
- local_model_data = st.session_state['local_models'][model_path]
464
- model_type = local_model_data.get('type')
465
- if model_type not in ['vision', 'ocr']: return f"Error: Loaded model '{os.path.basename(model_path)}' is not a Vision/OCR type."
466
- status_placeholder.warning(f"Local {model_type} Model ({os.path.basename(model_path)}): Processing logic depends on specific model. Placeholder active.")
467
- # --- ADD SPECIFIC LOCAL VISION/OCR MODEL LOGIC HERE ---
468
- # This section needs code tailored to the loaded model's processor/generate methods
469
- # Example placeholder:
470
- processor = local_model_data.get('processor')
471
- model = local_model_data.get('model')
472
- if processor and model:
473
- result_text = f"[Local {model_type} model processing needs implementation for {os.path.basename(model_path)}. Prompt: '{prompt}']"
474
- else:
475
- result_text = f"Error: Missing model or processor for local {model_type} model {os.path.basename(model_path)}."
476
- # --- END OF PLACEHOLDER ---
477
-
478
- elapsed = int(time.time() - start_time)
479
- status_placeholder.success(f"Image processing attempt completed in {elapsed}s.")
480
- return result_text
481
-
482
- async def process_hf_ocr(image: Image.Image, output_file: str, use_api: bool) -> str:
483
- """ Performs OCR using the process_image_hf function framework. """
484
- # Simple prompt for OCR task
485
- ocr_prompt = "Extract text content from this image."
486
- result = process_image_hf(image, ocr_prompt, use_api=use_api) # Pass use_api flag
487
-
488
- # Save the result if it looks like text (basic check)
489
- if result and not result.startswith("Error") and not result.startswith("["):
490
- try:
491
- async with aiofiles.open(output_file, "w", encoding='utf-8') as f:
492
- await f.write(result)
493
- logger.info(f"HF OCR result saved to {output_file}")
494
- except IOError as e:
495
- logger.error(f"Failed to save HF OCR output to {output_file}: {e}")
496
- result += f"\n[Error saving file: {e}]" # Append error to result if save fails
497
-
498
- # --- CORRECTED BLOCK ---
499
- elif os.path.exists(output_file):
500
- # Remove file if processing failed or was just a placeholder message
501
- try:
502
- os.remove(output_file)
503
- except OSError:
504
- # Log error or just ignore if removal fails
505
- logger.warning(f"Could not remove potentially empty/failed OCR file: {output_file}")
506
- pass # Ignore removal error
507
- except Exception as e_rem: # Catch any other error during removal
508
- logger.warning(f"Error removing OCR file {output_file}: {e_rem}")
509
- pass
510
- # --- END CORRECTION ---
511
-
512
- return result
513
-
514
- # --- Character Functions (Keep from previous) -----------
515
- def randomize_character_content():
516
- 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..."]
517
- 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...'"]
518
- 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)
519
- return name, gender, intro, greeting
520
-
521
- def save_character(character_data):
522
- characters = st.session_state.get('characters', []);
523
- if any(c['name'] == character_data['name'] for c in characters): st.error(f"Character name '{character_data['name']}' already exists."); return False
524
- characters.append(character_data); st.session_state['characters'] = characters
525
- try:
526
- 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
527
- except IOError as e: logger.error(f"Failed to save characters.json: {e}"); st.error(f"Failed to save character file: {e}"); return False
528
-
529
- def load_characters():
530
- if not os.path.exists("characters.json"): st.session_state['characters'] = []; return
531
- try:
532
- with open("characters.json", "r", encoding='utf-8') as f: characters = json.load(f)
533
- if isinstance(characters, list): st.session_state['characters'] = characters; logger.info(f"Loaded {len(characters)} characters.")
534
- else: st.session_state['characters'] = []; logger.warning("characters.json is not a list, resetting."); os.remove("characters.json")
535
- except (json.JSONDecodeError, IOError) as e:
536
- 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'] = []
537
- 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")
538
- except Exception as backup_e: logger.error(f"Could not backup corrupted character file: {backup_e}")
539
-
540
- # --- Utility: Clean stems (Keep from previous) ----------
541
- def clean_stem(fn: str) -> str:
542
- name = os.path.splitext(os.path.basename(fn))[0]; name = name.replace('-', ' ').replace('_', ' ')
543
- return name.strip().title()
544
-
545
- # --- PDF Creation Functions ---
546
- # Original image-only PDF function (might be removed or kept as an option)
547
- def make_image_sized_pdf(sources):
548
- # ... (kept same as previous version for now) ...
549
- if not sources: st.warning("No image sources provided for PDF generation."); return None
550
- buf = io.BytesIO(); c = canvas.Canvas(buf, pagesize=letter)
551
- try:
552
- for idx, src in enumerate(sources, start=1):
553
- status_placeholder = st.empty(); status_placeholder.info(f"Adding page {idx}/{len(sources)}: {os.path.basename(str(src))}...")
554
- try:
555
- filename = f'page_{idx}'
556
- if isinstance(src, str):
557
- 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
558
- img_obj = Image.open(src); filename = os.path.basename(src)
559
- elif hasattr(src, 'name'): # Handle uploaded file object
560
- src.seek(0); img_obj = Image.open(src); filename = getattr(src, 'name', f'uploaded_image_{idx}'); src.seek(0)
561
- else: continue # Skip unknown source type
562
- with img_obj:
563
- iw, ih = img_obj.size
564
- 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
565
- cap_h = 30; pw, ph = iw, ih + cap_h; c.setPageSize((pw, ph)); img_reader = ImageReader(img_obj)
566
- c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto')
567
- caption = clean_stem(filename); c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, caption)
568
- c.setFont('Helvetica', 8); c.setFillColorRGB(0.5, 0.5, 0.5); c.drawRightString(pw - 10, 8, f"Page {idx}")
569
- c.showPage(); status_placeholder.success(f"Added page {idx}/{len(sources)}: {filename}")
570
- 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}")
571
- except Exception as e: logger.error(f"Unexpected error adding page {idx} ({src}): {e}"); status_placeholder.error(f"Unexpected error on page {idx}: {e}")
572
- c.save(); buf.seek(0)
573
- if buf.getbuffer().nbytes < 100: st.error("PDF generation resulted in an empty file."); return None
574
- return buf.getvalue()
575
- except Exception as e: logger.error(f"Fatal error during PDF generation: {e}"); st.error(f"PDF Generation Failed: {e}"); return None
576
-
577
- # --- NEW Combined PDF Generation Function ---
578
- def make_combined_pdf(ordered_sources_info: List[Dict]) -> Optional[bytes]:
579
- if not ordered_sources_info:
580
- st.warning("No items selected for combined PDF generation.")
581
- return None
582
-
583
- buf = io.BytesIO()
584
- c = canvas.Canvas(buf, pagesize=letter)
585
- styles = getSampleStyleSheet()
586
- total_pages_generated = 0
587
-
588
- # Add page number function
589
- def draw_page_number(canvas, page_num, page_width, page_height):
590
- canvas.saveState()
591
- canvas.setFont('Helvetica', 8)
592
- canvas.setFillColorRGB(0.5, 0.5, 0.5)
593
- canvas.drawRightString(page_width - inch/2, inch/2, f"Page {page_num}")
594
- canvas.restoreState()
595
-
596
- for idx, item_info in enumerate(ordered_sources_info):
597
- filepath = item_info.get('filepath')
598
- file_type = item_info.get('type')
599
- filename = item_info.get('filename', f"item_{idx+1}")
600
- item_caption = clean_stem(filename)
601
-
602
- if not filepath: logger.warning(f"Skipping item {idx+1} due to missing filepath."); continue
603
- is_file_object = not isinstance(filepath, str)
604
- status_placeholder = st.empty()
605
- status_placeholder.info(f"Processing item {idx+1}/{len(ordered_sources_info)}: {filename} ({file_type})...")
606
-
607
- try:
608
- # --- IMAGE Processing ---
609
- if file_type == 'Image':
610
- if is_file_object: filepath.seek(0)
611
- try:
612
- img_obj = Image.open(filepath)
613
- with img_obj:
614
- iw, ih = img_obj.size
615
- if iw <= 0 or ih <= 0: raise ValueError("Invalid image dimensions")
616
- cap_h = 30; pw, ph = iw, ih + cap_h
617
- c.setPageSize((pw, ph)); img_reader = ImageReader(img_obj)
618
- c.drawImage(img_reader, 0, cap_h, width=iw, height=ih, preserveAspectRatio=True, anchor='c', mask='auto')
619
- c.setFont('Helvetica', 12); c.setFillColorRGB(0, 0, 0); c.drawCentredString(pw / 2, cap_h / 2 + 3, item_caption)
620
- total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
621
- c.showPage()
622
- finally:
623
- if is_file_object: filepath.seek(0)
624
-
625
- # --- PDF Processing ---
626
- elif file_type == 'PDF':
627
- src_doc = None
628
- try:
629
- if is_file_object: filepath.seek(0); pdf_bytes = filepath.read(); src_doc = fitz.open("pdf", pdf_bytes)
630
- else: src_doc = fitz.open(filepath)
631
- if len(src_doc) == 0: st.warning(f"Skipping empty PDF: {filename}"); continue
632
- for i, page in enumerate(src_doc):
633
- page_rect = page.rect; pw, ph = page_rect.width, page_rect.height
634
- if pw <= 0 or ph <= 0: continue
635
- c.setPageSize((pw, ph))
636
- pix = page.get_pixmap(dpi=150) # Render as image
637
- if pix.width > 0 and pix.height > 0:
638
- img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); img_reader = ImageReader(img)
639
- c.drawImage(img_reader, 0, 0, width=pw, height=ph)
640
- else: c.setFont('Helvetica', 10); c.setFillColorRGB(1,0,0); c.drawCentredString(pw/2, ph/2, f"Failed to render page {i+1} preview")
641
- overlay_text = f"{item_caption} (p{i+1})"; c.setFont('Helvetica', 8); c.setFillColorRGB(0, 0, 0, alpha=0.6); c.drawString(10, 10, overlay_text)
642
- total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
643
- c.showPage()
644
- finally:
645
- if src_doc: src_doc.close()
646
- if is_file_object: filepath.seek(0)
647
-
648
- # --- TEXT/MARKDOWN Processing ---
649
- elif file_type == 'Text':
650
- if is_file_object:
651
- filepath.seek(0)
652
- try: text_content = filepath.read().decode('utf-8')
653
- except: text_content = filepath.read().decode('latin-1', errors='replace')
654
- else:
655
- with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: text_content = f.read()
656
-
657
- temp_buf = io.BytesIO()
658
- temp_doc = SimpleDocTemplate(temp_buf, pagesize=letter, leftMargin=inch, rightMargin=inch, topMargin=inch, bottomMargin=inch)
659
- story = [Paragraph(f"Content from: {item_caption}", styles['h2']), Spacer(1, 0.2*inch)]
660
- # Use Preformatted for simple text dump
661
- story.append(Preformatted(text_content, styles['Code']))
662
- temp_doc.build(story)
663
- temp_buf.seek(0)
664
-
665
- text_pdf = fitz.open("pdf", temp_buf.read())
666
- for i, page in enumerate(text_pdf):
667
- page_rect = page.rect; pw, ph = page_rect.width, page_rect.height
668
- c.setPageSize((pw, ph)); pix = page.get_pixmap(dpi=150)
669
- if pix.width > 0 and pix.height > 0:
670
- img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); img_reader = ImageReader(img)
671
- c.drawImage(img_reader, 0, 0, width=pw, height=ph)
672
- else: c.setFont('Helvetica', 10); c.setFillColorRGB(1,0,0); c.drawCentredString(pw/2, ph/2, f"Failed to render text page {i+1}")
673
- total_pages_generated += 1; draw_page_number(c, total_pages_generated, pw, ph)
674
- c.showPage()
675
- text_pdf.close()
676
-
677
- else: # Unknown type
678
- logger.warning(f"Unsupported file type for PDF combination: {filename} ({file_type})")
679
- c.setPageSize(letter); c.setFont('Helvetica-Bold', 14); c.setFillColorRGB(0.7, 0.7, 0); c.drawCentredString(letter[0] / 2, letter[1] / 2 + 20, f"Unsupported File: {filename}")
680
- c.setFont('Helvetica', 10); c.drawCentredString(letter[0] / 2, letter[1] / 2 - 20, f"Type: {file_type}. Cannot include.")
681
- total_pages_generated += 1; draw_page_number(c, total_pages_generated, letter[0], letter[1])
682
- c.showPage()
683
-
684
- except Exception as item_err:
685
- logger.error(f"Error processing item {filename} for PDF: {item_err}", exc_info=True)
686
- try: # Add error page
687
- c.setPageSize(letter); c.setFont('Helvetica-Bold', 14); c.setFillColorRGB(1, 0, 0); c.drawCentredString(letter[0] / 2, letter[1] / 2 + 20, f"Error processing: {filename}")
688
- c.setFont('Helvetica', 10); c.drawCentredString(letter[0] / 2, letter[1] / 2 - 20, f"{str(item_err)[:100]}"); total_pages_generated += 1; draw_page_number(c, total_pages_generated, letter[0], letter[1]); c.showPage()
689
- except: logger.error(f"Failed to add error page for {filename}")
690
- finally:
691
- status_placeholder.empty()
692
-
693
- if total_pages_generated == 0: st.error("No pages were successfully added."); return None
694
- try:
695
- c.save(); buf.seek(0)
696
- if buf.getbuffer().nbytes < 100: st.error("Combined PDF generation resulted empty."); return None
697
- return buf.getvalue()
698
- except Exception as e: logger.error(f"Fatal error during final PDF save: {e}"); st.error(f"PDF Save Failed: {e}"); return None
699
-
700
-
701
- # --- Sidebar Gallery Update Function (MODIFIED for Sort, PDF Preview Fix, Delete Fix) ---
702
- def get_sort_key(filename):
703
- ext = os.path.splitext(filename)[1].lower()
704
- if ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']: priority = 1
705
- elif ext in ['.md', '.txt']: priority = 2
706
- elif ext == '.pdf': priority = 3
707
- else: priority = 4
708
- return (priority, filename.lower())
709
-
710
- def update_gallery():
711
- st.sidebar.markdown("### Asset Gallery 📸📖")
712
- all_files_unsorted = get_gallery_files()
713
- all_files = sorted(all_files_unsorted, key=get_sort_key) # Apply sorting
714
-
715
- if not all_files: st.sidebar.info("No assets found."); return
716
- st.sidebar.caption(f"Found {len(all_files)} assets:")
717
-
718
- for idx, file in enumerate(all_files):
719
- st.session_state['unique_counter'] += 1
720
- unique_id = st.session_state['unique_counter']
721
- item_key_base = f"gallery_item_{os.path.basename(file)}_{unique_id}"
722
- basename = os.path.basename(file)
723
- st.sidebar.markdown(f"**{basename}**")
724
-
725
- try:
726
- file_ext = os.path.splitext(file)[1].lower()
727
- preview_failed = False
728
- # Previews with better error handling
729
- if file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']:
730
- try:
731
- with st.sidebar.expander("Preview", expanded=False): st.image(Image.open(file), use_container_width=True)
732
- except Exception as img_err: st.sidebar.warning(f"Img preview failed: {img_err}"); preview_failed = True
733
- elif file_ext == '.pdf':
734
- try:
735
- with st.sidebar.expander("Preview (Page 1)", expanded=False):
736
- doc = fitz.open(file)
737
- if len(doc) > 0:
738
- pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
739
- if pix.width > 0 and pix.height > 0: img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); st.image(img, use_container_width=True)
740
- else: st.warning("Failed to render PDF page."); preview_failed = True
741
- else: st.warning("Empty PDF")
742
- doc.close()
743
- except Exception as pdf_err: st.sidebar.warning(f"PDF preview failed: {pdf_err}"); logger.warning(f"PDF preview error {file}: {pdf_err}"); preview_failed = True
744
- elif file_ext in ['.md', '.txt']:
745
- try:
746
- with st.sidebar.expander("Preview (Start)", expanded=False):
747
- with open(file, 'r', encoding='utf-8', errors='ignore') as f: content_preview = f.read(200)
748
- st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text')
749
- except Exception as txt_err: st.sidebar.warning(f"Text preview failed: {txt_err}"); preview_failed = True
750
-
751
- # Actions
752
- action_cols = st.sidebar.columns(3)
753
- with action_cols[0]:
754
- checkbox_key = f"cb_{item_key_base}"
755
- st.session_state.setdefault('asset_checkboxes', {})
756
- st.session_state['asset_checkboxes'][file] = st.checkbox("Select", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
757
- with action_cols[1]:
758
- mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'}
759
- mime_type = mime_map.get(file_ext, "application/octet-stream"); dl_key = f"dl_{item_key_base}"
760
- try:
761
- with open(file, "rb") as fp: st.download_button(label="📥", data=fp, file_name=basename, mime=mime_type, key=dl_key, help="Download")
762
- except Exception as dl_e: st.error(f"DL Err: {dl_e}")
763
- with action_cols[2]:
764
- delete_key = f"del_{item_key_base}"
765
- if st.button("🗑️", key=delete_key, help=f"Delete {basename}"):
766
- delete_success = False
767
- try:
768
- os.remove(file)
769
- st.session_state['asset_checkboxes'].pop(file, None)
770
- if file in st.session_state.get('layout_snapshots', []): st.session_state['layout_snapshots'].remove(file) # Remove if also in old list
771
- logger.info(f"Deleted asset: {file}")
772
- st.toast(f"Deleted {basename}!", icon="✅")
773
- delete_success = True
774
- except OSError as e: logger.error(f"Error deleting file {file}: {e}"); st.error(f"Could not delete {basename}: {e}")
775
- except Exception as e: logger.error(f"Unexpected error deleting file {file}: {e}"); st.error(f"Could not delete {basename}: {e}")
776
- # Rerun to refresh the gallery list after attempting delete
777
- st.rerun()
778
-
779
- except FileNotFoundError: st.sidebar.error(f"File vanished: {basename}"); st.session_state['asset_checkboxes'].pop(file, None)
780
- except Exception as e: st.sidebar.error(f"Display Error: {basename}"); logger.error(f"Error displaying asset {file}: {e}")
781
- st.sidebar.markdown("---")
782
-
783
- # --- UI Elements -----------------------------------------
784
- # Sidebar Structure
785
- st.sidebar.subheader("🤖 Hugging Face Settings")
786
- # ... (HF API, Local Model, Params Expanders - code unchanged) ...
787
- with st.sidebar.expander("API Inference Settings", expanded=False):
788
- 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.")
789
- 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}")
790
- providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
791
- 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.")
792
- 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.")
793
- 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.")
794
- effective_api_model = st.session_state.hf_custom_api_model.strip() or st.session_state.hf_selected_api_model
795
- 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.")
796
- st.caption(f"Effective API Model: {effective_api_model}")
797
- with st.sidebar.expander("Local Model Selection", expanded=True):
798
- if not _transformers_available: st.warning("Transformers library not found.")
799
- else:
800
- local_model_options = ["None"] + list(st.session_state.get('local_models', {}).keys())
801
- current_selection = st.session_state.get('selected_local_model_path'); current_selection = current_selection if current_selection in local_model_options else "None"
802
- 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 loaded local model.")
803
- st.session_state.selected_local_model_path = selected_path if selected_path != "None" else None
804
- if st.session_state.selected_local_model_path:
805
- model_info = st.session_state.local_models[st.session_state.selected_local_model_path]
806
- st.caption(f"Type: {model_info.get('type', '?')} | Device: {model_info.get('model').device if model_info.get('model') else 'N/A'}")
807
- else: st.caption("No local model selected.")
808
- with st.sidebar.expander("Generation Parameters", expanded=False):
809
- 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")
810
- 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")
811
- 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")
812
- 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.")
813
- 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.")
814
-
815
- st.sidebar.markdown("---")
816
- # Gallery is rendered later by calling update_gallery()
817
-
818
- # --- App Title & Main Area ---
819
- st.title("Vision & Layout Titans (HF) 🚀🖼️📄")
820
- st.markdown("Combined App: PDF Layout Generator + Hugging Face Powered AI Tools")
821
-
822
- # Warning for missing libraries in main area if sidebar not ready
823
- if not _transformers_available:
824
- st.warning("AI/ML libraries (torch, transformers) not found. Local model features disabled.")
825
- elif not _diffusers_available:
826
- st.warning("Diffusers library not found. Diffusion model features disabled.")
827
-
828
-
829
- # --- Main Application Tabs ---
830
- tabs_to_create = [
831
- "Combined PDF Generator 📄", # Renamed Tab 0
832
- "Camera Snap 📷",
833
- "Download PDFs 📥",
834
- "Build Titan (Local Models) 🌱",
835
- "PDF Page Process (HF) 📄", # Clarified name
836
- "Image Process (HF) 🖼️",
837
- "Text Process (HF) 📝",
838
- "Test OCR (HF) 🔍",
839
- "Test Image Gen (Diffusers) 🎨",
840
- "Character Editor 🧑‍🎨",
841
- "Character Gallery 🖼️",
842
- ]
843
- tabs = st.tabs(tabs_to_create)
844
-
845
- # --- Tab Implementations ---
846
-
847
- # --- Tab 1: Combined PDF Generator (OVERHAULED) ---
848
- with tabs[0]:
849
- st.header("Combined PDF Generator 📄➕🖼️➕...")
850
- st.markdown("Select assets (Images, PDFs, Text/MD) from the sidebar gallery, reorder them, and generate a combined PDF.")
851
-
852
- # --- Get Selected Files ---
853
- selected_files_paths = [
854
- f for f, selected in st.session_state.get('asset_checkboxes', {}).items()
855
- if selected and os.path.exists(f) # Ensure file still exists
856
- ]
857
-
858
- if not selected_files_paths:
859
- st.info("👈 Select one or more assets from the sidebar gallery using the checkboxes.")
860
- else:
861
- st.info(f"{len(selected_files_paths)} assets selected from gallery.")
862
-
863
- # --- Populate DataFrame for Reordering ---
864
- combined_records = []
865
- for idx, filepath in enumerate(selected_files_paths):
866
- filename = os.path.basename(filepath)
867
- ext = os.path.splitext(filename)[1].lower()
868
- file_type = "Unknown"
869
- if ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff']: file_type = "Image"
870
- elif ext == '.pdf': file_type = "PDF"
871
- elif ext in ['.md', '.txt']: file_type = "Text"
872
-
873
- combined_records.append({
874
- "filename": filename,
875
- "filepath": filepath, # Keep the path
876
- "type": file_type,
877
- "order": idx, # Initial order based on selection
878
- })
879
-
880
- combined_df_initial = pd.DataFrame(combined_records)
881
-
882
- st.markdown("#### Reorder Selected Assets for PDF")
883
- st.caption("Edit the 'Order' column or drag rows to set the sequence for the combined PDF.")
884
-
885
- edited_combined_df = st.data_editor(
886
- combined_df_initial,
887
- column_config={
888
- "filename": st.column_config.TextColumn("Filename", disabled=True),
889
- "filepath": None, # Hide filepath column
890
- "type": st.column_config.TextColumn("Type", disabled=True),
891
- "order": st.column_config.NumberColumn(
892
- "Order",
893
- min_value=0,
894
- # max_value=len(combined_df_initial)-1, # Max can cause issues if rows added/removed by user selection change
895
- step=1,
896
- required=True,
897
- ),
898
- },
899
- hide_index=True,
900
- use_container_width=True,
901
- num_rows="dynamic", # Allow drag-and-drop reordering
902
- key="combined_pdf_editor"
903
- )
904
-
905
- # Sort by the edited 'order' column
906
- ordered_combined_df = edited_combined_df.sort_values('order').reset_index(drop=True)
907
-
908
- # Prepare list of dicts for the PDF generation function
909
- ordered_sources_info_for_pdf = ordered_combined_df[['filepath', 'type', 'filename']].to_dict('records')
910
-
911
- # --- Generate & Download ---
912
- st.subheader("Generate Combined PDF")
913
- if st.button("🖋️ Generate Combined PDF", key="generate_combined_pdf_btn"):
914
- if not ordered_sources_info_for_pdf:
915
- st.warning("No items available after reordering.")
916
- else:
917
- with st.spinner("Generating combined PDF... This might take a while."):
918
- combined_pdf_bytes = make_combined_pdf(ordered_sources_info_for_pdf)
919
-
920
- if combined_pdf_bytes:
921
- # Create filename
922
- now = datetime.now(pytz.timezone("US/Central"))
923
- prefix = now.strftime("%Y%m%d-%H%M%p")
924
- first_item_name = clean_stem(ordered_sources_info_for_pdf[0].get('filename','combined'))
925
- combined_pdf_fname = f"{prefix}_Combined_{first_item_name}.pdf"
926
- combined_pdf_fname = re.sub(r'[^\w\-\.\_]', '_', combined_pdf_fname) # Sanitize
927
-
928
- st.success(f"✅ Combined PDF ready: **{combined_pdf_fname}**")
929
- st.download_button(
930
- "⬇️ Download Combined PDF",
931
- data=combined_pdf_bytes,
932
- file_name=combined_pdf_fname,
933
- mime="application/pdf",
934
- key="download_combined_pdf_btn"
935
- )
936
- # Add preview (optional, might be slow for large combined PDFs)
937
- # ... (preview logic similar to other tabs if desired) ...
938
- else:
939
- st.error("Combined PDF generation failed. Check logs or input files.")
940
-
941
-
942
- # --- Tab 2: Camera Snap ---
943
- with tabs[1]:
944
- st.header("Camera Snap 📷")
945
- st.subheader("Single Capture (Adds to General Gallery)")
946
- cols = st.columns(2)
947
- with cols[0]:
948
- cam0_img = st.camera_input("Take a picture - Cam 0", key="main_cam0")
949
- if cam0_img:
950
- filename = generate_filename("cam0_snap");
951
- if st.session_state.get('cam0_file') and os.path.exists(st.session_state['cam0_file']):
952
- try:
953
- os.remove(st.session_state['cam0_file'])
954
- except OSError:
955
- pass
956
- try:
957
- with open(filename, "wb") as f: f.write(cam0_img.getvalue())
958
- 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}")
959
- update_gallery(); # Refresh sidebar without rerun
960
- 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}")
961
- with cols[1]:
962
- cam1_img = st.camera_input("Take a picture - Cam 1", key="main_cam1")
963
- if cam1_img:
964
- filename = generate_filename("cam1_snap")
965
- if st.session_state.get('cam1_file') and os.path.exists(st.session_state['cam1_file']):
966
- try:
967
- os.remove(st.session_state['cam1_file'])
968
- except OSError:
969
- pass
970
- try:
971
- with open(filename, "wb") as f: f.write(cam1_img.getvalue())
972
- 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}")
973
- update_gallery(); # Refresh sidebar without rerun
974
- 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}")
975
-
976
-
977
- # --- Tab 3: Download PDFs ---
978
- with tabs[2]:
979
- st.header("Download PDFs 📥")
980
- st.markdown("Download PDFs from URLs and optionally create image snapshots.")
981
- if st.button("Load Example arXiv URLs 📚", key="load_examples"):
982
- 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/"]
983
- st.session_state['pdf_urls_input'] = "\n".join(example_urls)
984
- 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")
985
- if st.button("Robo-Download PDFs 🤖", key="download_pdfs_button"):
986
- urls = [url.strip() for url in url_input.strip().split("\n") if url.strip()]
987
- if not urls: st.warning("Please enter at least one URL.")
988
- else:
989
- progress_bar = st.progress(0); status_text = st.empty(); total_urls = len(urls); download_count = 0; existing_pdfs = get_pdf_files()
990
- for idx, url in enumerate(urls):
991
- 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)
992
- if os.path.exists(output_path): # Check existence properly
993
- st.info(f"Already exists: {os.path.basename(output_path)}")
994
- st.session_state['downloaded_pdfs'][url] = output_path
995
- # Ensure checkbox state is preserved or reset if needed
996
- st.session_state['asset_checkboxes'][output_path] = st.session_state['asset_checkboxes'].get(output_path, False)
997
- else:
998
- if download_pdf(url, output_path):
999
- 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)
1000
- else: st.error(f"Failed to download: {url}")
1001
- status_text.success(f"Download process complete! Successfully downloaded {download_count} new PDFs.")
1002
- if download_count > 0: update_gallery(); # Refresh sidebar without rerun
1003
-
1004
- st.subheader("Create Snapshots from Gallery PDFs")
1005
- 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")
1006
- 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}
1007
- 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"}
1008
- resolution = resolution_map[snapshot_mode]; mode_key = mode_key_map[snapshot_mode]
1009
- if st.button("Snapshot Selected PDFs 📸", key="snapshot_selected_pdfs"):
1010
- selected_pdfs = [path for path in get_gallery_files(['pdf']) if st.session_state['asset_checkboxes'].get(path, False)]
1011
- if not selected_pdfs: st.warning("No PDFs selected in the sidebar gallery!")
1012
- else:
1013
- st.info(f"Starting snapshot process for {len(selected_pdfs)} selected PDF(s)..."); snapshot_count = 0; total_snapshots_generated = 0
1014
- for pdf_path in selected_pdfs:
1015
- if not os.path.exists(pdf_path): st.warning(f"File not found: {pdf_path}. Skipping."); continue
1016
- new_snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key, resolution))
1017
- if new_snapshots:
1018
- snapshot_count += 1; total_snapshots_generated += len(new_snapshots)
1019
- st.write(f"Snapshots for {os.path.basename(pdf_path)}:"); cols = st.columns(3)
1020
- for i, snap_path in enumerate(new_snapshots):
1021
- with cols[i % 3]:
1022
- try: st.image(Image.open(snap_path), caption=os.path.basename(snap_path), use_container_width=True)
1023
- except Exception as snap_img_err: st.warning(f"Cannot display snap {os.path.basename(snap_path)}: {snap_img_err}")
1024
- st.session_state['asset_checkboxes'][snap_path] = False # Add to gallery
1025
- if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery(); # Refresh sidebar without rerun
1026
- else: st.warning("No snapshots were generated. Check logs or PDF files.")
1027
-
1028
-
1029
- # --- Tab 4: Build Titan (Local Models) ---
1030
- with tabs[3]:
1031
- st.header("Build Titan (Local Models) 🌱")
1032
- st.markdown("Download and save models from Hugging Face Hub for local use.")
1033
- if not _transformers_available:
1034
- st.error("Transformers library not available. Cannot download or load local models.")
1035
- else:
1036
- build_model_type = st.selectbox("Select Model Type", ["Causal LM", "Vision/Multimodal", "OCR", "Diffusion"], key="build_type_local")
1037
- st.subheader(f"Download {build_model_type} Model")
1038
- 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")
1039
- 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.")
1040
- 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).")
1041
-
1042
- if st.button(f"Download & Save '{hf_model_id}' Locally", key="build_download_button", disabled=not hf_model_id or not local_model_name):
1043
- local_name_check = re.sub(r'[^\w\-]+', '_', local_model_name) # Sanitize proposed name for path check
1044
- potential_path_base = os.path.join(f"{build_model_type.split('/')[0].lower()}_models", local_name_check)
1045
-
1046
- if any(os.path.basename(p) == local_name_check for p in get_local_model_paths(build_model_type.split('/')[0].lower())):
1047
- st.error(f"A local model folder named '{local_name_check}' already exists. Choose a different local name.")
1048
- else:
1049
- model_type_map = {"Causal LM": "causal", "Vision/Multimodal": "vision", "OCR": "ocr", "Diffusion": "diffusion"}
1050
- model_type_short = model_type_map.get(build_model_type, "unknown")
1051
- config = LocalModelConfig(name=local_model_name, hf_id=hf_model_id, model_type=model_type_short)
1052
- save_path = config.get_full_path()
1053
- os.makedirs(os.path.dirname(save_path), exist_ok=True)
1054
- st.info(f"Attempting to download '{hf_model_id}' to '{save_path}'..."); progress_bar_build = st.progress(0); status_text_build = st.empty()
1055
- token_build = st.session_state.hf_custom_key or HF_TOKEN or None
1056
- try:
1057
- if build_model_type == "Diffusion":
1058
- if not _diffusers_available: raise ImportError("Diffusers library required.")
1059
- status_text_build.text("Downloading diffusion 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)
1060
- st.session_state.local_models[save_path] = {'type': 'diffusion', 'hf_id': hf_model_id, 'model':None, 'processor':None} # Mark as downloaded
1061
- st.success(f"Diffusion model '{hf_model_id}' downloaded and saved to {save_path}")
1062
- del pipeline_obj # Free memory
1063
- else:
1064
- status_text_build.text("Downloading model components...")
1065
- if model_type_short == 'causal': model_class, proc_tok_class = AutoModelForCausalLM, AutoTokenizer; proc_name="tokenizer"
1066
- elif model_type_short == 'vision': model_class, proc_tok_class = AutoModelForVision2Seq, AutoProcessor; proc_name="processor"
1067
- elif model_type_short == 'ocr': model_class, proc_tok_class = AutoModelForVision2Seq, AutoProcessor; proc_name="processor"
1068
- else: raise ValueError(f"Unknown model type: {model_type_short}")
1069
-
1070
- model_obj = model_class.from_pretrained(hf_model_id, token=token_build); model_obj.save_pretrained(save_path)
1071
- status_text_build.text(f"Model saved. Downloading {proc_name}..."); proc_tok_obj = proc_tok_class.from_pretrained(hf_model_id, token=token_build); proc_tok_obj.save_pretrained(save_path)
1072
- status_text_build.text(f"Components saved. Loading '{local_model_name}' into memory...")
1073
- device = "cuda" if torch.cuda.is_available() else "cpu"
1074
- # Use trust_remote_code cautiously if needed for specific models
1075
- reloaded_model = model_class.from_pretrained(save_path).to(device)
1076
- reloaded_proc_tok = proc_tok_class.from_pretrained(save_path)
1077
- st.session_state.local_models[save_path] = {'type': model_type_short, 'hf_id': hf_model_id, 'model': reloaded_model, proc_name: reloaded_proc_tok}
1078
- # Add tokenizer specifically if it's nested in processor
1079
- if proc_name == "processor" and hasattr(reloaded_proc_tok, 'tokenizer'):
1080
- st.session_state.local_models[save_path]['tokenizer'] = reloaded_proc_tok.tokenizer
1081
- st.success(f"{build_model_type} model '{hf_model_id}' downloaded to {save_path} and loaded ({device})."); st.session_state.selected_local_model_path = save_path
1082
- del model_obj, proc_tok_obj # Free memory from download cache if possible
1083
- except (RepositoryNotFoundError, GatedRepoError) as e: st.error(f"Download failed: Repo not found or requires access/token. Error: {e}"); logger.error(f"Download failed for {hf_model_id}: {e}"); #if os.path.exists(save_path): shutil.rmtree(save_path)
1084
- except ImportError as e: st.error(f"Download failed: Library missing. {e}"); logger.error(f"ImportError for {hf_model_id}: {e}")
1085
- except Exception as e: st.error(f"Download error: {e}"); logger.error(f"Download failed for {hf_model_id}: {e}", exc_info=True); #if os.path.exists(save_path): shutil.rmtree(save_path)
1086
- finally: progress_bar_build.progress(1.0); status_text_build.empty(); #st.rerun() # Rerun removed
1087
-
1088
- st.subheader("Manage Local Models")
1089
- # Refresh list for display
1090
- loaded_model_paths = list(st.session_state.get('local_models', {}).keys())
1091
- if not loaded_model_paths: st.info("No models downloaded yet.")
1092
- else:
1093
- models_df_data = []
1094
- for path in loaded_model_paths:
1095
- data = st.session_state.local_models.get(path, {}) # Safely get data
1096
- models_df_data.append({
1097
- "Local Name": os.path.basename(path), "Type": data.get('type', '?'),
1098
- "HF ID": data.get('hf_id', '?'), "Loaded": "Yes" if data.get('model') else "No", "Path": path })
1099
- 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"])
1100
- model_to_delete = st.selectbox("Select model to delete", [""] + [os.path.basename(p) for p in loaded_model_paths], key="delete_model_select")
1101
- if model_to_delete and st.button(f"Delete Local Model '{model_to_delete}'", type="primary"):
1102
- path_to_delete = next((p for p in loaded_model_paths if os.path.basename(p) == model_to_delete), None)
1103
- if path_to_delete:
1104
- try:
1105
- # Explicitly delete model objects from memory first if they exist
1106
- if path_to_delete in st.session_state.local_models:
1107
- model_data_to_del = st.session_state.local_models[path_to_delete]
1108
- if model_data_to_del.get('model'): del model_data_to_del['model']
1109
- if model_data_to_del.get('tokenizer'): del model_data_to_del['tokenizer']
1110
- if model_data_to_del.get('processor'): del model_data_to_del['processor']
1111
- if _transformers_available and torch.cuda.is_available(): torch.cuda.empty_cache() # Try to clear VRAM
1112
-
1113
- # Remove from session state
1114
- st.session_state.local_models.pop(path_to_delete, None)
1115
- if st.session_state.selected_local_model_path == path_to_delete: st.session_state.selected_local_model_path = None
1116
- # Delete from disk
1117
- if os.path.exists(path_to_delete): shutil.rmtree(path_to_delete)
1118
- st.success(f"Deleted model '{model_to_delete}'."); logger.info(f"Deleted local model: {path_to_delete}"); st.rerun()
1119
- 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}")
1120
-
1121
-
1122
- # --- Tab 5: PDF Process (HF) ---
1123
- with tabs[4]:
1124
- st.header("PDF Page Process with HF Models 📄")
1125
- st.markdown("Upload PDFs, view pages, and extract text/info using selected HF models (API or Local Vision/OCR).")
1126
- 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 for vision/OCR).")
1127
- 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} (likely image-to-text)")
1128
- else:
1129
- if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
1130
- else: st.warning("No local model selected.")
1131
-
1132
- uploaded_pdfs_process_hf = st.file_uploader("Upload PDF files to process", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader_hf")
1133
- if uploaded_pdfs_process_hf:
1134
- process_all_pages_pdf = st.checkbox("Process All Pages (can be slow/expensive)", value=False, key="pdf_process_all_hf")
1135
- pdf_prompt = st.text_area("Prompt for PDF Page Processing", "Extract the text content from this page.", key="pdf_process_prompt_hf")
1136
- if st.button("Process Uploaded PDFs with HF", key="process_uploaded_pdfs_hf"):
1137
- 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.")
1138
- else:
1139
- 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()
1140
- for pdf_file in uploaded_pdfs_process_hf:
1141
- output_placeholder_hf.markdown(f"--- \n### Processing: {pdf_file.name}")
1142
- try:
1143
- pdf_bytes = pdf_file.read(); doc = fitz.open("pdf", pdf_bytes); num_pages = len(doc)
1144
- pages_to_process = range(num_pages) if process_all_pages_pdf else range(min(1, num_pages))
1145
- output_placeholder_hf.info(f"Processing {len(pages_to_process)} of {num_pages} pages..."); doc_text = f"## File: {pdf_file.name}\n\n"
1146
- for i in pages_to_process:
1147
- page_placeholder = output_placeholder_hf.empty(); page_placeholder.info(f"Processing Page {i + 1}/{num_pages}...")
1148
- page = doc.load_page(i); pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
1149
- cols_pdf = output_placeholder_hf.columns(2); cols_pdf[0].image(img, caption=f"Page {i+1}", use_container_width=True)
1150
- with cols_pdf[1], st.spinner("Processing page with HF model..."): hf_text = process_image_hf(img, pdf_prompt, use_api=pdf_use_api)
1151
- st.text_area(f"Result (Page {i+1})", hf_text, height=250, key=f"pdf_hf_out_{pdf_file.name}_{i}")
1152
- doc_text += f"### Page {i + 1}\n\n{hf_text}\n\n---\n\n"; total_pages_processed_hf += 1; page_placeholder.empty()
1153
- combined_text_output_hf += doc_text; doc.close()
1154
- except Exception as e: output_placeholder_hf.error(f"Error processing {pdf_file.name}: {str(e)}")
1155
- if total_pages_processed_hf > 0:
1156
- st.markdown("--- \n### Combined Processing Results"); st.text_area("Full Output", combined_text_output_hf, height=400, key="combined_pdf_hf_output")
1157
- output_filename_pdf_hf = generate_filename("hf_processed_pdfs", "md")
1158
- try:
1159
- with open(output_filename_pdf_hf, "w", encoding="utf-8") as f: f.write(combined_text_output_hf)
1160
- st.success(f"Combined output saved to {output_filename_pdf_hf}")
1161
- st.markdown(get_download_link(output_filename_pdf_hf, "text/markdown", "Download Combined MD"), unsafe_allow_html=True)
1162
- st.session_state['asset_checkboxes'][output_filename_pdf_hf] = False; update_gallery() # Refresh sidebar
1163
- except IOError as e: st.error(f"Failed to save combined output file: {e}")
1164
-
1165
- # --- Tab 6: Image Process (HF) ---
1166
- with tabs[5]:
1167
- st.header("Image Process with HF Models 🖼️")
1168
- st.markdown("Upload images and process them using selected HF models (API or Local).")
1169
- img_use_api = st.radio("Choose Processing Method", ["Hugging Face API", "Loaded Local Model"], key="img_process_source_hf", horizontal=True)
1170
- 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} (likely image-to-text)")
1171
- else:
1172
- if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
1173
- else: st.warning("No local model selected.")
1174
- img_prompt_hf = st.text_area("Prompt for Image Processing", "Describe this image in detail.", key="img_process_prompt_hf")
1175
- uploaded_images_process_hf = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader_hf")
1176
- if uploaded_images_process_hf:
1177
- if st.button("Process Uploaded Images with HF", key="process_images_hf"):
1178
- 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.")
1179
- else:
1180
- 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()
1181
- for img_file in uploaded_images_process_hf:
1182
- output_img_placeholder_hf.markdown(f"### Processing: {img_file.name}")
1183
- try:
1184
- 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)
1185
- with cols_img_hf[1], st.spinner("Processing image with HF model..."): hf_img_text = process_image_hf(img, img_prompt_hf, use_api=img_use_api)
1186
- st.text_area(f"Result", hf_img_text, height=300, key=f"img_hf_out_{img_file.name}")
1187
- combined_img_text_hf += f"## Image: {img_file.name}\n\n{hf_img_text}\n\n---\n\n"; images_processed_hf += 1
1188
- except UnidentifiedImageError: output_img_placeholder_hf.error(f"Invalid Image: {img_file.name}. Skipping.")
1189
- except Exception as e: output_img_placeholder_hf.error(f"Error processing {img_file.name}: {str(e)}")
1190
- if images_processed_hf > 0:
1191
- st.markdown("--- \n### Combined Processing Results"); st.text_area("Full Output", combined_img_text_hf, height=400, key="combined_img_hf_output")
1192
- output_filename_img_hf = generate_filename("hf_processed_images", "md")
1193
- try:
1194
- with open(output_filename_img_hf, "w", encoding="utf-8") as f: f.write(combined_img_text_hf)
1195
- 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)
1196
- st.session_state['asset_checkboxes'][output_filename_img_hf] = False; update_gallery() # Refresh sidebar
1197
- except IOError as e: st.error(f"Failed to save combined output file: {e}")
1198
-
1199
- # --- Tab 7: Text Process (HF) ---
1200
- with tabs[6]:
1201
- st.header("Text Process with HF Models 📝")
1202
- st.markdown("Process Markdown (.md) or Text (.txt) files using selected HF models (API or Local).")
1203
- text_use_api = st.radio("Choose Processing Method", ["Hugging Face API", "Loaded Local Model"], key="text_process_source_hf", horizontal=True)
1204
- 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}")
1205
- else:
1206
- if st.session_state.selected_local_model_path: st.info(f"Using Local Model: {os.path.basename(st.session_state.selected_local_model_path)}")
1207
- else: st.warning("No local model selected.")
1208
- text_files_hf = get_gallery_files(['md', 'txt'])
1209
- if not text_files_hf: st.warning("No .md or .txt files in gallery to process.")
1210
- else:
1211
- 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")
1212
- if selected_text_file_hf:
1213
- st.write(f"Selected: {os.path.basename(selected_text_file_hf)}")
1214
- try:
1215
- with open(selected_text_file_hf, "r", encoding="utf-8", errors='ignore') as f: content_text_hf = f.read()
1216
- st.text_area("File Content Preview", content_text_hf[:1000] + ("..." if len(content_text_hf) > 1000 else ""), height=200, key="text_content_preview_hf")
1217
- prompt_text_hf = st.text_area("Enter Prompt for this File", "Summarize the key points of this text.", key="text_individual_prompt_hf")
1218
- if st.button(f"Process '{os.path.basename(selected_text_file_hf)}' with HF", key=f"process_text_hf_btn"):
1219
- 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.")
1220
- else:
1221
- 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)
1222
- st.markdown("### Processing Result"); st.markdown(result_text_processed)
1223
- output_filename_text_hf = generate_filename(f"hf_processed_{os.path.splitext(os.path.basename(selected_text_file_hf))[0]}", "md")
1224
- try:
1225
- with open(output_filename_text_hf, "w", encoding="utf-8") as f: f.write(result_text_processed)
1226
- 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)
1227
- st.session_state['asset_checkboxes'][output_filename_text_hf] = False; update_gallery() # Refresh sidebar
1228
- except IOError as e: st.error(f"Failed to save result file: {e}")
1229
- except FileNotFoundError: st.error("Selected file not found.")
1230
- except Exception as e: st.error(f"Error reading file: {e}")
1231
-
1232
- # --- Tab 8: Test OCR (HF) ---
1233
- with tabs[7]:
1234
- st.header("Test OCR with HF Models 🔍")
1235
- st.markdown("Select an image/PDF and run OCR using HF models (API or Local - requires suitable local model).")
1236
- 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.")
1237
- if ocr_use_api == "Loaded Local OCR Model":
1238
- if st.session_state.selected_local_model_path:
1239
- model_info = st.session_state.local_models.get(st.session_state.selected_local_model_path,{})
1240
- model_type = model_info.get('type'); model_name = os.path.basename(st.session_state.selected_local_model_path)
1241
- if model_type != 'ocr': st.warning(f"Selected model ({model_name}) is type '{model_type}', not 'ocr'. Results may be poor.")
1242
- else: st.info(f"Using Local OCR Model: {model_name}")
1243
- else: st.warning("No local model selected.")
1244
-
1245
- gallery_files_ocr_hf = get_gallery_files(['png', 'jpg', 'jpeg', 'pdf'])
1246
- if not gallery_files_ocr_hf: st.warning("No images or PDFs in gallery.")
1247
- else:
1248
- 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")
1249
- if selected_file_ocr_hf:
1250
- 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 = ""
1251
- try:
1252
- if file_ext_ocr_hf in ['.png', '.jpg', '.jpeg']: image_to_ocr_hf = Image.open(selected_file_ocr_hf)
1253
- elif file_ext_ocr_hf == '.pdf':
1254
- doc = fitz.open(selected_file_ocr_hf)
1255
- 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)"
1256
- else: st.warning("Selected PDF is empty.")
1257
- doc.close()
1258
- if image_to_ocr_hf:
1259
- st.image(image_to_ocr_hf, caption=f"Image for OCR{page_info_hf}", use_container_width=True)
1260
- if st.button("Run HF OCR on this Image 🚀", key="ocr_run_button_hf"):
1261
- 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.")
1262
- else:
1263
- 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
1264
- 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))
1265
- st.session_state['processing']['ocr'] = False; st.text_area("OCR Result", ocr_result_hf, height=300, key="ocr_result_display_hf")
1266
- if ocr_result_hf and not ocr_result_hf.startswith("Error") and not ocr_result_hf.startswith("["):
1267
- entry = f"HF OCR: {selected_file_ocr_hf}{page_info_hf} -> {output_ocr_file_hf}"
1268
- st.session_state['history'].append(entry)
1269
- if len(ocr_result_hf) > 5: 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() # Refresh sidebar
1270
- else: st.warning("OCR output seems short/empty.")
1271
- else: st.error(f"OCR failed. {ocr_result_hf}")
1272
- except Exception as e: st.error(f"Error loading file for OCR: {e}")
1273
-
1274
- # --- Tab 9: Test Image Gen (Diffusers) ---
1275
- with tabs[8]:
1276
- st.header("Test Image Generation (Diffusers) 🎨")
1277
- st.markdown("Generate images using Stable Diffusion models loaded locally via the Diffusers library.")
1278
- if not _diffusers_available: st.error("Diffusers library is required.")
1279
- else:
1280
- local_diffusion_paths = get_local_model_paths("diffusion") # Check diffusion_models folder
1281
- if not local_diffusion_paths: st.warning("No local diffusion models found. Download one using the 'Build Titan' tab."); selected_diffusion_model_path = None
1282
- 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")
1283
- prompt_imggen_diff = st.text_area("Image Generation Prompt", "A photorealistic cat wearing sunglasses, studio lighting", key="imggen_prompt_diff")
1284
- neg_prompt_imggen_diff = st.text_area("Negative Prompt (Optional)", "ugly, deformed, blurry, low quality", key="imggen_neg_prompt_diff")
1285
- 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")
1286
- if st.button("Generate Image 🚀", key="imggen_run_button_diff", disabled=not selected_diffusion_model_path):
1287
- if not prompt_imggen_diff: st.warning("Please enter a prompt.")
1288
- else:
1289
- status_imggen = st.empty()
1290
- try:
1291
- status_imggen.info(f"Loading diffusion pipeline: {os.path.basename(selected_diffusion_model_path)}..."); device = "cuda" if _transformers_available and torch.cuda.is_available() else "cpu"; dtype = torch.float16 if device == "cuda" else torch.float32
1292
- pipe = StableDiffusionPipeline.from_pretrained(selected_diffusion_model_path, torch_dtype=dtype).to(device); pipe.safety_checker = None # Optional
1293
- status_imggen.info(f"Generating image on {device} ({dtype})..."); start_gen_time = time.time()
1294
- gen_output = pipe(prompt=prompt_imggen_diff, negative_prompt=neg_prompt_imggen_diff or None, num_inference_steps=steps_imggen_diff, guidance_scale=guidance_imggen_diff)
1295
- gen_image = gen_output.images[0]; elapsed_gen = int(time.time() - start_gen_time); status_imggen.success(f"Image generated in {elapsed_gen}s!")
1296
- output_imggen_file_diff = generate_filename("diffusion_gen", "png"); gen_image.save(output_imggen_file_diff)
1297
- st.image(gen_image, caption=f"Generated: {output_imggen_file_diff}", use_container_width=True)
1298
- st.markdown(get_download_link(output_imggen_file_diff, "image/png", "Download Generated Image"), unsafe_allow_html=True)
1299
- st.session_state['asset_checkboxes'][output_imggen_file_diff] = False; update_gallery() # Refresh sidebar
1300
- st.session_state['history'].append(f"Diffusion Gen: '{prompt_imggen_diff[:30]}...' -> {output_imggen_file_diff}")
1301
- except ImportError: st.error("Diffusers or Torch library not found.")
1302
- except Exception as e: st.error(f"Image generation failed: {e}"); logger.error(f"Diffusion generation failed for {selected_diffusion_model_path}: {e}", exc_info=True)
1303
- finally:
1304
- if 'pipe' in locals():
1305
- del pipe; torch.cuda.empty_cache() if device == "cuda" else None # Clear VRAM
1306
-
1307
- # --- Tab 10: Character Editor ---
1308
- with tabs[9]:
1309
- st.header("Character Editor 🧑‍🎨"); st.subheader("Create Your Character")
1310
- load_characters(); existing_char_names = [c['name'] for c in st.session_state.get('characters', [])]
1311
- form_key = f"character_form_{st.session_state.get('char_form_reset_key', 0)}"
1312
- with st.form(key=form_key):
1313
- st.markdown("**Create New Character**")
1314
- if st.form_submit_button("Randomize Content 🎲"): st.session_state['char_form_reset_key'] += 1; st.rerun()
1315
- rand_name, rand_gender, rand_intro, rand_greeting = randomize_character_content()
1316
- name_char = st.text_input("Name (3-25 chars...)", value=rand_name, max_chars=25, key="char_name_input")
1317
- gender_char = st.radio("Gender", ["Male", "Female"], index=["Male", "Female"].index(rand_gender), key="char_gender_radio")
1318
- intro_char = st.text_area("Intro (Public description)", value=rand_intro, max_chars=300, height=100, key="char_intro_area")
1319
- greeting_char = st.text_area("Greeting (First message)", value=rand_greeting, max_chars=300, height=100, key="char_greeting_area")
1320
- tags_char = st.text_input("Tags (comma-separated)", "OC, friendly", key="char_tags_input")
1321
- submitted = st.form_submit_button("Create Character ✨")
1322
- if submitted:
1323
- error = False; # Validation checks...
1324
- if not (3 <= len(name_char) <= 25): st.error("Name must be 3-25 characters."); error = True
1325
- if not re.match(r'^[a-zA-Z0-9 _-]+$', name_char): st.error("Name contains invalid characters."); error = True
1326
- if name_char in existing_char_names: st.error(f"Name '{name_char}' already exists!"); error = True
1327
- if not intro_char or not greeting_char: st.error("Intro/Greeting cannot be empty."); error = True
1328
- if not error:
1329
- tag_list = [tag.strip() for tag in tags_char.split(',') if tag.strip()]
1330
- 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}
1331
- if save_character(character_data): st.success(f"Character '{name_char}' created!"); st.session_state['char_form_reset_key'] += 1; st.rerun()
1332
-
1333
- # --- Tab 11: Character Gallery ---
1334
- with tabs[10]:
1335
- st.header("Character Gallery 🖼️"); load_characters(); characters_list = st.session_state.get('characters', [])
1336
- if not characters_list: st.warning("No characters created yet.")
1337
- else:
1338
- st.subheader(f"Your Characters ({len(characters_list)})"); search_term = st.text_input("Search Characters by Name", key="char_gallery_search")
1339
- if search_term: characters_list = [c for c in characters_list if search_term.lower() in c['name'].lower()]
1340
- cols_char_gallery = st.columns(3); chars_to_delete = []
1341
- for idx, char in enumerate(characters_list):
1342
- with cols_char_gallery[idx % 3], st.container(border=True):
1343
- st.markdown(f"**{char['name']}**"); st.caption(f"Gender: {char.get('gender', 'N/A')}")
1344
- st.markdown("**Intro:**"); st.markdown(f"> {char.get('intro', '')}")
1345
- st.markdown("**Greeting:**"); st.markdown(f"> {char.get('greeting', '')}")
1346
- st.caption(f"Tags: {', '.join(char.get('tags', ['N/A']))}"); st.caption(f"Created: {char.get('created_at', 'N/A')}")
1347
- delete_key_char = f"delete_char_{char['name']}_{idx}";
1348
- if st.button(f"Delete", key=delete_key_char, type="primary", help=f"Delete {char['name']}"): chars_to_delete.append(char['name']) # Shorten button label
1349
- if chars_to_delete:
1350
- current_characters = st.session_state.get('characters', []); updated_characters = [c for c in current_characters if c['name'] not in chars_to_delete]
1351
- st.session_state['characters'] = updated_characters
1352
- try:
1353
- with open("characters.json", "w", encoding='utf-8') as f: json.dump(updated_characters, f, indent=2)
1354
- logger.info(f"Deleted characters: {', '.join(chars_to_delete)}"); st.success(f"Deleted: {', '.join(chars_to_delete)}"); st.rerun()
1355
- except IOError as e: logger.error(f"Failed to save characters.json after deletion: {e}"); st.error("Failed to update character file.")
1356
-
1357
- # --- Footer and Persistent Sidebar Elements ------------
1358
- st.sidebar.markdown("---")
1359
- # Update Sidebar Gallery (Call this at the end to reflect all changes)
1360
- update_gallery()
1361
-
1362
- # Action Logs in Sidebar
1363
- st.sidebar.subheader("Action Logs 📜")
1364
- log_expander = st.sidebar.expander("View Logs", expanded=False)
1365
- with log_expander:
1366
- # Display logs in reverse order (newest first)
1367
- log_text = "\n".join([f"{record.levelname}: {record.message}" for record in reversed(log_records)])
1368
- st.code(log_text, language='log')
1369
-
1370
- # History in Sidebar
1371
- st.sidebar.subheader("Session History 📜")
1372
- history_expander = st.sidebar.expander("View History", expanded=False)
1373
- with history_expander:
1374
- for entry in reversed(st.session_state.get("history", [])):
1375
- if entry: history_expander.write(f"- {entry}")
1376
-
1377
- st.sidebar.markdown("---")
1378
- st.sidebar.info("Using Hugging Face models for AI tasks.")
1379
- st.sidebar.caption("App Modified by AI Assistant")
 
1
+ # app.py (streamlined and modularized for clarity and maintainability)
2
+ import streamlit as st
3
  import os
 
 
4
  import glob
5
+ import base64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  import pandas as pd
7
+ import fitz
8
+ from PIL import Image
9
+ from io import BytesIO
10
+ from datetime import datetime
11
  from reportlab.pdfgen import canvas
12
  from reportlab.lib.utils import ImageReader
13
+ from markdown2 import markdown
 
 
 
 
 
 
 
 
 
14
 
15
+ # --- Config ---
16
  st.set_page_config(
17
+ page_title="Vision & Layout Titans 🚀",
18
  page_icon="🤖",
19
+ layout="wide"
 
 
 
 
 
 
20
  )
21
 
22
+ # --- Helper Functions ---
23
+ def get_files(exts):
24
+ files = []
25
+ for ext in exts:
26
+ files.extend(glob.glob(f'*.{ext}'))
27
+ return sorted([f for f in files if f.lower() != "readme.md"])
28
+
29
+ def image_to_pdf(images, md_files):
30
+ buffer = BytesIO()
31
+ c = canvas.Canvas(buffer)
32
+
33
+ for img_path in images:
34
+ img = Image.open(img_path)
35
+ width, height = img.size
36
+ c.setPageSize((width, height))
37
+ c.drawImage(ImageReader(img), 0, 0, width, height)
38
+ c.showPage()
39
+
40
+ for md_path in md_files:
41
+ with open(md_path, 'r', encoding='utf-8') as f:
42
+ md_content = f.read()
43
+ html = markdown(md_content)
44
+ c.setPageSize((595, 842)) # A4 size
45
+ c.setFont("Helvetica", 10)
46
+ c.drawString(50, 800, md_content[:1000]) # Simplified render
47
+ c.showPage()
48
+
49
+ c.save()
50
+ buffer.seek(0)
51
+ return buffer
52
+
53
+ def render_pdf_gallery(pdf_files):
54
+ for pdf_file in pdf_files:
55
+ doc = fitz.open(pdf_file)
56
+ page = doc.load_page(0)
57
+ pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
58
+ img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
59
+ st.image(img, caption=os.path.basename(pdf_file))
60
+
61
+ # --- Sidebar Management ---
62
+ st.sidebar.header("Content Management")
63
+
64
+ image_files = get_files(['png', 'jpg', 'jpeg'])
65
+ pdf_files = get_files(['pdf'])
66
+ md_files = get_files(['md'])
67
+
68
+ selected_images = st.sidebar.multiselect("Select Images", image_files)
69
+ selected_md_files = st.sidebar.multiselect("Select Markdown Files", md_files)
70
+
71
+ # PDF Gallery
72
+ if st.sidebar.checkbox("Show PDF Gallery"):
73
+ render_pdf_gallery(pdf_files)
74
+
75
+ # PDF Generation with rearrangement
76
+ st.sidebar.subheader("Reorder Content for PDF")
77
+ content_df = pd.DataFrame({
78
+ "Content": selected_images + selected_md_files,
79
+ "Type": ["Image"] * len(selected_images) + ["Markdown"] * len(selected_md_files),
80
+ "Order": range(len(selected_images) + len(selected_md_files))
81
+ })
82
+
83
+ edited_df = st.sidebar.data_editor(content_df, use_container_width=True)
84
+ sorted_contents = edited_df.sort_values('Order')['Content'].tolist()
85
+
86
+ if st.sidebar.button("Generate PDF"):
87
+ sorted_images = [item for item in sorted_contents if item in selected_images]
88
+ sorted_md_files = [item for item in sorted_contents if item in selected_md_files]
89
+ pdf_buffer = image_to_pdf(sorted_images, sorted_md_files)
90
+ st.sidebar.download_button("Download PDF", pdf_buffer, "output.pdf")
91
+
92
+ # Deletion
93
+ st.sidebar.subheader("Delete Files")
94
+ file_to_delete = st.sidebar.selectbox("Select file to delete", image_files + pdf_files + md_files)
95
+ if st.sidebar.button("Delete Selected File"):
96
+ os.remove(file_to_delete)
97
+ st.sidebar.success(f"Deleted {file_to_delete}")
98
+ st.rerun()
99
+
100
+ # --- Main Page ---
101
+ st.title("Vision & Layout Titans 🚀")
102
+ st.markdown("### Manage, View, and Export Your Files Easily!")
103
+
104
+ # Display selected images
105
+ st.subheader("Selected Images")
106
+ for img_path in selected_images:
107
+ img = Image.open(img_path)
108
+ st.image(img, caption=os.path.basename(img_path))
109
+
110
+ # Display selected markdown
111
+ st.subheader("Selected Markdown")
112
+ for md_path in selected_md_files:
113
+ with open(md_path, 'r', encoding='utf-8') as f:
114
+ md_content = f.read()
115
+ st.markdown(md_content[:500] + '...')