awacke1 commited on
Commit
d6e5038
·
verified ·
1 Parent(s): 457db8f

Create 033025-1.app.py

Browse files
Files changed (1) hide show
  1. 033025-1.app.py +682 -0
033025-1.app.py ADDED
@@ -0,0 +1,682 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import aiofiles
2
+ import asyncio
3
+ import base64
4
+ import fitz
5
+ import glob
6
+ import logging
7
+ import os
8
+ import pandas as pd
9
+ import pytz
10
+ import random
11
+ import re
12
+ import requests
13
+ import shutil
14
+ import streamlit as st
15
+ import time
16
+ import torch
17
+ import zipfile
18
+
19
+ from dataclasses import dataclass
20
+ from datetime import datetime
21
+ from diffusers import StableDiffusionPipeline
22
+ from io import BytesIO
23
+ from openai import OpenAI
24
+ from PIL import Image
25
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
26
+ from typing import Optional
27
+
28
+ # 🤖 OpenAI wizardry: Summon your API magic!
29
+ client = OpenAI(
30
+ api_key=os.getenv('OPENAI_API_KEY'),
31
+ organization=os.getenv('OPENAI_ORG_ID')
32
+ )
33
+
34
+ # 📜 Logging activated: Capturing chaos and calm!
35
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
36
+ logger = logging.getLogger(__name__)
37
+ log_records = []
38
+ class LogCaptureHandler(logging.Handler):
39
+ def emit(self, record):
40
+ log_records.append(record)
41
+ logger.addHandler(LogCaptureHandler())
42
+
43
+ # 🎨 Streamlit styling: Designing a cosmic interface!
44
+ st.set_page_config(
45
+ page_title="AI Vision & SFT Titans 🚀",
46
+ page_icon="🤖",
47
+ layout="wide",
48
+ initial_sidebar_state="expanded",
49
+ menu_items={
50
+ 'Get Help': 'https://huggingface.co/awacke1',
51
+ 'Report a Bug': 'https://huggingface.co/spaces/awacke1',
52
+ 'About': "AI Vision & SFT Titans: PDFs, OCR, Image Gen, Line Drawings, Custom Diffusion, and SFT on CPU! 🌌"
53
+ }
54
+ )
55
+
56
+ # Set up default session state values.
57
+ st.session_state.setdefault('history', []) # History: starting fresh if empty!
58
+ st.session_state.setdefault('builder', None) # Builder: set up if missing.
59
+ st.session_state.setdefault('model_loaded', False) # Model Loaded: not loaded by default.
60
+ st.session_state.setdefault('processing', {}) # Processing: initialize as an empty dict.
61
+ st.session_state.setdefault('asset_checkboxes', {}) # Asset Checkboxes: default to an empty dict.
62
+ st.session_state.setdefault('downloaded_pdfs', {}) # Downloaded PDFs: start with none.
63
+ st.session_state.setdefault('unique_counter', 0) # Unique Counter: initialize to zero.
64
+ st.session_state.setdefault('selected_model_type', "Causal LM")
65
+ st.session_state.setdefault('selected_model', "None")
66
+ st.session_state.setdefault('cam0_file', None)
67
+ st.session_state.setdefault('cam1_file', None)
68
+
69
+ # Create a single container for the asset gallery in the sidebar.
70
+ if 'asset_gallery_container' not in st.session_state:
71
+ st.session_state['asset_gallery_container'] = st.sidebar.empty()
72
+
73
+ @dataclass # ModelConfig: A blueprint for model configurations.
74
+ class ModelConfig:
75
+ name: str
76
+ base_model: str
77
+ size: str
78
+ domain: Optional[str] = None
79
+ model_type: str = "causal_lm"
80
+ @property
81
+ def model_path(self):
82
+ return f"models/{self.name}"
83
+
84
+ @dataclass # DiffusionConfig: Where diffusion magic takes shape.
85
+ class DiffusionConfig:
86
+ name: str
87
+ base_model: str
88
+ size: str
89
+ domain: Optional[str] = None
90
+ @property
91
+ def model_path(self):
92
+ return f"diffusion_models/{self.name}"
93
+
94
+ class ModelBuilder:
95
+ def __init__(self):
96
+ self.config = None
97
+ self.model = None
98
+ self.tokenizer = None
99
+ self.jokes = [
100
+ "Why did the AI go to therapy? Too many layers to unpack! 😂",
101
+ "Training complete! Time for a binary coffee break. ☕",
102
+ "I told my neural network a joke; it couldn't stop dropping bits! 🤖",
103
+ "I asked the AI for a pun, and it said, 'I'm punning on parallel processing!' 😄",
104
+ "Debugging my code is like a stand-up routine—always a series of exceptions! 😆"
105
+ ]
106
+ def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
107
+ with st.spinner(f"Loading {model_path}... ⏳"):
108
+ self.model = AutoModelForCausalLM.from_pretrained(model_path)
109
+ self.tokenizer = AutoTokenizer.from_pretrained(model_path)
110
+ if self.tokenizer.pad_token is None:
111
+ self.tokenizer.pad_token = self.tokenizer.eos_token
112
+ if config:
113
+ self.config = config
114
+ self.model.to("cuda" if torch.cuda.is_available() else "cpu")
115
+ st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
116
+ return self
117
+ def save_model(self, path: str):
118
+ with st.spinner("Saving model... 💾"):
119
+ os.makedirs(os.path.dirname(path), exist_ok=True)
120
+ self.model.save_pretrained(path)
121
+ self.tokenizer.save_pretrained(path)
122
+ st.success(f"Model saved at {path}! ✅")
123
+
124
+ class DiffusionBuilder:
125
+ def __init__(self):
126
+ self.config = None
127
+ self.pipeline = None
128
+ def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
129
+ with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
130
+ self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
131
+ if config:
132
+ self.config = config
133
+ st.success("Diffusion model loaded! 🎨")
134
+ return self
135
+ def save_model(self, path: str):
136
+ with st.spinner("Saving diffusion model... 💾"):
137
+ os.makedirs(os.path.dirname(path), exist_ok=True)
138
+ self.pipeline.save_pretrained(path)
139
+ st.success(f"Diffusion model saved at {path}! ✅")
140
+ def generate(self, prompt: str):
141
+ return self.pipeline(prompt, num_inference_steps=20).images[0]
142
+
143
+ def generate_filename(sequence, ext="png"):
144
+ return f"{sequence}_{time.strftime('%d%m%Y%H%M%S')}.{ext}"
145
+
146
+ def pdf_url_to_filename(url):
147
+ return re.sub(r'[<>:"/\\|?*]', '_', url) + ".pdf"
148
+
149
+ def get_download_link(file_path, mime_type="application/pdf", label="Download"):
150
+ return f'<a href="data:{mime_type};base64,{base64.b64encode(open(file_path, "rb").read()).decode()}" download="{os.path.basename(file_path)}">{label}</a>'
151
+
152
+ def zip_directory(directory_path, zip_path):
153
+ with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
154
+ [zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))
155
+ for root, _, files in os.walk(directory_path) for file in files]
156
+
157
+ def get_model_files(model_type="causal_lm"):
158
+ return [d for d in glob.glob("models/*" if model_type == "causal_lm" else "diffusion_models/*") if os.path.isdir(d)] or ["None"]
159
+
160
+ def get_gallery_files(file_types=["png", "pdf"]):
161
+ return sorted(list({f for ext in file_types for f in glob.glob(f"*.{ext}")}))
162
+
163
+ def get_pdf_files():
164
+ return sorted(glob.glob("*.pdf"))
165
+
166
+ def download_pdf(url, output_path):
167
+ try:
168
+ response = requests.get(url, stream=True, timeout=10)
169
+ if response.status_code == 200:
170
+ with open(output_path, "wb") as f:
171
+ for chunk in response.iter_content(chunk_size=8192):
172
+ f.write(chunk)
173
+ ret = True
174
+ else:
175
+ ret = False
176
+ except requests.RequestException as e:
177
+ logger.error(f"Failed to download {url}: {e}")
178
+ ret = False
179
+ return ret
180
+
181
+ # Async PDF Snapshot: Snap your PDF pages without blocking.
182
+ async def process_pdf_snapshot(pdf_path, mode="single"):
183
+ start_time = time.time()
184
+ status = st.empty()
185
+ status.text(f"Processing PDF Snapshot ({mode})... (0s)")
186
+ try:
187
+ doc = fitz.open(pdf_path)
188
+ output_files = []
189
+ if mode == "single":
190
+ page = doc[0]
191
+ pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
192
+ output_file = generate_filename("single", "png")
193
+ pix.save(output_file)
194
+ output_files.append(output_file)
195
+ elif mode == "twopage":
196
+ for i in range(min(2, len(doc))):
197
+ page = doc[i]
198
+ pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
199
+ output_file = generate_filename(f"twopage_{i}", "png")
200
+ pix.save(output_file)
201
+ output_files.append(output_file)
202
+ elif mode == "allpages":
203
+ for i in range(len(doc)):
204
+ page = doc[i]
205
+ pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
206
+ output_file = generate_filename(f"page_{i}", "png")
207
+ pix.save(output_file)
208
+ output_files.append(output_file)
209
+ doc.close()
210
+ elapsed = int(time.time() - start_time)
211
+ status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
212
+ return output_files
213
+ except Exception as e:
214
+ status.error(f"Failed to process PDF: {str(e)}")
215
+ return []
216
+
217
+ # Async OCR: Convert images to text.
218
+ async def process_ocr(image, output_file):
219
+ start_time = time.time()
220
+ status = st.empty()
221
+ status.text("Processing GOT-OCR2_0... (0s)")
222
+ tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
223
+ model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
224
+ temp_file = f"temp_{int(time.time())}.png"
225
+ image.save(temp_file)
226
+ result = model.chat(tokenizer, temp_file, ocr_type='ocr')
227
+ os.remove(temp_file)
228
+ elapsed = int(time.time() - start_time)
229
+ status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
230
+ async with aiofiles.open(output_file, "w") as f:
231
+ await f.write(result)
232
+ return result
233
+
234
+ # Async Image Gen: Your image genie.
235
+ async def process_image_gen(prompt, output_file):
236
+ start_time = time.time()
237
+ status = st.empty()
238
+ status.text("Processing Image Gen... (0s)")
239
+ pipeline = (st.session_state['builder'].pipeline
240
+ if st.session_state.get('builder') and isinstance(st.session_state['builder'], DiffusionBuilder)
241
+ and st.session_state['builder'].pipeline
242
+ else StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu"))
243
+ gen_image = pipeline(prompt, num_inference_steps=20).images[0]
244
+ elapsed = int(time.time() - start_time)
245
+ status.text(f"Image Gen completed in {elapsed}s!")
246
+ gen_image.save(output_file)
247
+ return gen_image
248
+
249
+ # GPT-Image Interpreter: Turning pixels into prose!
250
+ def process_image_with_prompt(image, prompt, model="gpt-4o-mini", detail="auto"):
251
+ buffered = BytesIO()
252
+ image.save(buffered, format="PNG")
253
+ img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
254
+ messages = [{
255
+ "role": "user",
256
+ "content": [
257
+ {"type": "text", "text": prompt},
258
+ {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}", "detail": detail}}
259
+ ]
260
+ }]
261
+ try:
262
+ response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
263
+ return response.choices[0].message.content
264
+ except Exception as e:
265
+ return f"Error processing image with GPT: {str(e)}"
266
+
267
+ # GPT-Text Alchemist: Merging prompt and text.
268
+ def process_text_with_prompt(text, prompt, model="gpt-4o-mini"):
269
+ messages = [{"role": "user", "content": f"{prompt}\n\n{text}"}]
270
+ try:
271
+ response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
272
+ return response.choices[0].message.content
273
+ except Exception as e:
274
+ return f"Error processing text with GPT: {str(e)}"
275
+
276
+ # ----------------- SIDEBAR UPDATES -----------------
277
+
278
+ # Sidebar: Gallery Settings
279
+ st.sidebar.subheader("Gallery Settings")
280
+ st.session_state.setdefault('gallery_size', 2)
281
+ st.session_state['gallery_size'] = st.sidebar.slider("Gallery Size", 1, 10, st.session_state['gallery_size'], key="gallery_size_slider")
282
+
283
+ # ----------------- TAB SETUP -----------------
284
+ tabs = st.tabs([
285
+ "Camera Snap 📷", "Download PDFs 📥", "Test OCR 🔍", "Build Titan 🌱",
286
+ "Test Image Gen 🎨", "PDF Process 📄", "Image Process 🖼️", "MD Gallery 📚"
287
+ ])
288
+ (tab_camera, tab_download, tab_ocr, tab_build, tab_imggen, tab_pdf_process, tab_image_process, tab_md_gallery) = tabs
289
+
290
+ # ----------------- TAB: Camera Snap -----------------
291
+ with tab_camera:
292
+ st.header("Camera Snap 📷")
293
+ st.subheader("Single Capture")
294
+ cols = st.columns(2)
295
+ with cols[0]:
296
+ cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
297
+ if cam0_img:
298
+ filename = generate_filename("cam0")
299
+ if st.session_state['cam0_file'] and os.path.exists(st.session_state['cam0_file']):
300
+ os.remove(st.session_state['cam0_file'])
301
+ with open(filename, "wb") as f:
302
+ f.write(cam0_img.getvalue())
303
+ st.session_state['cam0_file'] = filename
304
+ entry = f"Snapshot from Cam 0: {filename}"
305
+ st.session_state['history'].append(entry)
306
+ st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
307
+ logger.info(f"Saved snapshot from Camera 0: {filename}")
308
+ with cols[1]:
309
+ cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
310
+ if cam1_img:
311
+ filename = generate_filename("cam1")
312
+ if st.session_state['cam1_file'] and os.path.exists(st.session_state['cam1_file']):
313
+ os.remove(st.session_state['cam1_file'])
314
+ with open(filename, "wb") as f:
315
+ f.write(cam1_img.getvalue())
316
+ st.session_state['cam1_file'] = filename
317
+ entry = f"Snapshot from Cam 1: {filename}"
318
+ st.session_state['history'].append(entry)
319
+ st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
320
+ logger.info(f"Saved snapshot from Camera 1: {filename}")
321
+
322
+ # ----------------- TAB: Download PDFs -----------------
323
+ with tab_download:
324
+ st.header("Download PDFs 📥")
325
+ if st.button("Examples 📚"):
326
+ example_urls = [
327
+ "https://arxiv.org/pdf/2308.03892",
328
+ "https://arxiv.org/pdf/1912.01703",
329
+ "https://arxiv.org/pdf/2408.11039",
330
+ "https://arxiv.org/pdf/2109.10282",
331
+ "https://arxiv.org/pdf/2112.10752",
332
+ "https://arxiv.org/pdf/2308.11236",
333
+ "https://arxiv.org/pdf/1706.03762",
334
+ "https://arxiv.org/pdf/2006.11239",
335
+ "https://arxiv.org/pdf/2305.11207",
336
+ "https://arxiv.org/pdf/2106.09685",
337
+ "https://arxiv.org/pdf/2005.11401",
338
+ "https://arxiv.org/pdf/2106.10504"
339
+ ]
340
+ st.session_state['pdf_urls'] = "\n".join(example_urls)
341
+ url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200)
342
+ if st.button("Robo-Download 🤖"):
343
+ urls = url_input.strip().split("\n")
344
+ progress_bar = st.progress(0)
345
+ status_text = st.empty()
346
+ total_urls = len(urls)
347
+ existing_pdfs = get_pdf_files()
348
+ for idx, url in enumerate(urls):
349
+ if url:
350
+ output_path = pdf_url_to_filename(url)
351
+ status_text.text(f"Fetching {idx + 1}/{total_urls}: {os.path.basename(output_path)}...")
352
+ if output_path not in existing_pdfs:
353
+ if download_pdf(url, output_path):
354
+ st.session_state['downloaded_pdfs'][url] = output_path
355
+ logger.info(f"Downloaded PDF from {url} to {output_path}")
356
+ entry = f"Downloaded PDF: {output_path}"
357
+ st.session_state['history'].append(entry)
358
+ st.session_state['asset_checkboxes'][output_path] = True
359
+ else:
360
+ st.error(f"Failed to nab {url} 😿")
361
+ else:
362
+ st.info(f"Already got {os.path.basename(output_path)}! Skipping... 🐾")
363
+ st.session_state['downloaded_pdfs'][url] = output_path
364
+ progress_bar.progress((idx + 1) / total_urls)
365
+ status_text.text("Robo-Download complete! 🚀")
366
+ mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], key="download_mode")
367
+ if st.button("Snapshot Selected 📸"):
368
+ selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(path, False)]
369
+ if selected_pdfs:
370
+ for pdf_path in selected_pdfs:
371
+ if not os.path.exists(pdf_path):
372
+ st.warning(f"File not found: {pdf_path}. Skipping.")
373
+ continue
374
+ mode_key = {"Single Page (High-Res)": "single",
375
+ "Two Pages (High-Res)": "twopage",
376
+ "All Pages (High-Res)": "allpages"}[mode]
377
+ snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
378
+ for snapshot in snapshots:
379
+ st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
380
+ st.session_state['asset_checkboxes'][snapshot] = True
381
+ # No update_gallery() call here; will update once later.
382
+ else:
383
+ st.warning("No PDFs selected for snapshotting! Check some boxes in the sidebar.")
384
+
385
+ # ----------------- TAB: Test OCR -----------------
386
+ with tab_ocr:
387
+ st.header("Test OCR 🔍")
388
+ all_files = get_gallery_files()
389
+ if all_files:
390
+ if st.button("OCR All Assets 🚀"):
391
+ full_text = "# OCR Results\n\n"
392
+ for file in all_files:
393
+ if file.endswith('.png'):
394
+ image = Image.open(file)
395
+ else:
396
+ doc = fitz.open(file)
397
+ pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
398
+ image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
399
+ doc.close()
400
+ output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt")
401
+ result = asyncio.run(process_ocr(image, output_file))
402
+ full_text += f"## {os.path.basename(file)}\n\n{result}\n\n"
403
+ entry = f"OCR Test: {file} -> {output_file}"
404
+ st.session_state['history'].append(entry)
405
+ md_output_file = f"full_ocr_{int(time.time())}.md"
406
+ with open(md_output_file, "w") as f:
407
+ f.write(full_text)
408
+ st.success(f"Full OCR saved to {md_output_file}")
409
+ st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
410
+ selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
411
+ if selected_file:
412
+ if selected_file.endswith('.png'):
413
+ image = Image.open(selected_file)
414
+ else:
415
+ doc = fitz.open(selected_file)
416
+ pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
417
+ image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
418
+ doc.close()
419
+ st.image(image, caption="Input Image", use_container_width=True)
420
+ if st.button("Run OCR 🚀", key="ocr_run"):
421
+ output_file = generate_filename("ocr_output", "txt")
422
+ st.session_state['processing']['ocr'] = True
423
+ result = asyncio.run(process_ocr(image, output_file))
424
+ entry = f"OCR Test: {selected_file} -> {output_file}"
425
+ st.session_state['history'].append(entry)
426
+ st.text_area("OCR Result", result, height=200, key="ocr_result")
427
+ st.success(f"OCR output saved to {output_file}")
428
+ st.session_state['processing']['ocr'] = False
429
+ if selected_file.endswith('.pdf') and st.button("OCR All Pages 🚀", key="ocr_all_pages"):
430
+ doc = fitz.open(selected_file)
431
+ full_text = f"# OCR Results for {os.path.basename(selected_file)}\n\n"
432
+ for i in range(len(doc)):
433
+ pix = doc[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
434
+ image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
435
+ output_file = generate_filename(f"ocr_page_{i}", "txt")
436
+ result = asyncio.run(process_ocr(image, output_file))
437
+ full_text += f"## Page {i + 1}\n\n{result}\n\n"
438
+ entry = f"OCR Test: {selected_file} Page {i + 1} -> {output_file}"
439
+ st.session_state['history'].append(entry)
440
+ md_output_file = f"full_ocr_{os.path.basename(selected_file)}_{int(time.time())}.md"
441
+ with open(md_output_file, "w") as f:
442
+ f.write(full_text)
443
+ st.success(f"Full OCR saved to {md_output_file}")
444
+ st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
445
+ else:
446
+ st.warning("No assets in gallery yet. Use Camera Snap or Download PDFs!")
447
+
448
+ # ----------------- TAB: Build Titan -----------------
449
+ with tab_build:
450
+ st.header("Build Titan 🌱")
451
+ model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
452
+ base_model = st.selectbox(
453
+ "Select Tiny Model",
454
+ ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM"
455
+ else ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"]
456
+ )
457
+ model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
458
+ domain = st.text_input("Target Domain", "general")
459
+ if st.button("Download Model ⬇️"):
460
+ config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(
461
+ name=model_name, base_model=base_model, size="small", domain=domain
462
+ )
463
+ builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
464
+ builder.load_model(base_model, config)
465
+ builder.save_model(config.model_path)
466
+ st.session_state['builder'] = builder
467
+ st.session_state['model_loaded'] = True
468
+ st.session_state['selected_model_type'] = model_type
469
+ st.session_state['selected_model'] = config.model_path
470
+ entry = f"Built {model_type} model: {model_name}"
471
+ st.session_state['history'].append(entry)
472
+ st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
473
+ st.experimental_rerun()
474
+
475
+ # ----------------- TAB: Test Image Gen -----------------
476
+ with tab_imggen:
477
+ st.header("Test Image Gen 🎨")
478
+ all_files = get_gallery_files()
479
+ if all_files:
480
+ selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
481
+ if selected_file:
482
+ if selected_file.endswith('.png'):
483
+ image = Image.open(selected_file)
484
+ else:
485
+ doc = fitz.open(selected_file)
486
+ pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
487
+ image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
488
+ doc.close()
489
+ st.image(image, caption="Reference Image", use_container_width=True)
490
+ prompt = st.text_area("Prompt", "Generate a neon superhero version of this image", key="gen_prompt")
491
+ if st.button("Run Image Gen 🚀", key="gen_run"):
492
+ output_file = generate_filename("gen_output", "png")
493
+ st.session_state['processing']['gen'] = True
494
+ result = asyncio.run(process_image_gen(prompt, output_file))
495
+ entry = f"Image Gen Test: {prompt} -> {output_file}"
496
+ st.session_state['history'].append(entry)
497
+ st.image(result, caption="Generated Image", use_container_width=True)
498
+ st.success(f"Image saved to {output_file}")
499
+ st.session_state['processing']['gen'] = False
500
+ else:
501
+ st.warning("No images or PDFs in gallery yet. Use Camera Snap or Download PDFs!")
502
+
503
+ # ----------------- TAB: PDF Process -----------------
504
+ with tab_pdf_process:
505
+ st.header("PDF Process")
506
+ st.subheader("Upload PDFs for GPT-based text extraction")
507
+ gpt_models = ["gpt-4o", "gpt-4o-mini"]
508
+ selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="pdf_gpt_model")
509
+ detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="pdf_detail_level")
510
+ uploaded_pdfs = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader")
511
+ view_mode = st.selectbox("View Mode", ["Single Page", "Double Page"], key="pdf_view_mode")
512
+ if st.button("Process Uploaded PDFs", key="process_pdfs"):
513
+ combined_text = ""
514
+ for pdf_file in uploaded_pdfs:
515
+ pdf_bytes = pdf_file.read()
516
+ temp_pdf_path = f"temp_{pdf_file.name}"
517
+ with open(temp_pdf_path, "wb") as f:
518
+ f.write(pdf_bytes)
519
+ try:
520
+ doc = fitz.open(temp_pdf_path)
521
+ st.write(f"Processing {pdf_file.name} with {len(doc)} pages")
522
+ if view_mode == "Single Page":
523
+ for i, page in enumerate(doc):
524
+ pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
525
+ img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
526
+ st.image(img, caption=f"{pdf_file.name} Page {i+1}")
527
+ gpt_text = process_image_with_prompt(img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level)
528
+ combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"
529
+ else:
530
+ pages = list(doc)
531
+ for i in range(0, len(pages), 2):
532
+ if i+1 < len(pages):
533
+ pix1 = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
534
+ img1 = Image.frombytes("RGB", [pix1.width, pix1.height], pix1.samples)
535
+ pix2 = pages[i+1].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
536
+ img2 = Image.frombytes("RGB", [pix2.width, pix2.height], pix2.samples)
537
+ total_width = img1.width + img2.width
538
+ max_height = max(img1.height, img2.height)
539
+ combined_img = Image.new("RGB", (total_width, max_height))
540
+ combined_img.paste(img1, (0, 0))
541
+ combined_img.paste(img2, (img1.width, 0))
542
+ st.image(combined_img, caption=f"{pdf_file.name} Pages {i+1}-{i+2}")
543
+ gpt_text = process_image_with_prompt(combined_img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level)
544
+ combined_text += f"\n## {pdf_file.name} - Pages {i+1}-{i+2}\n\n{gpt_text}\n"
545
+ else:
546
+ pix = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
547
+ img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
548
+ st.image(img, caption=f"{pdf_file.name} Page {i+1}")
549
+ gpt_text = process_image_with_prompt(img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level)
550
+ combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"
551
+ doc.close()
552
+ except Exception as e:
553
+ st.error(f"Error processing {pdf_file.name}: {str(e)}")
554
+ finally:
555
+ os.remove(temp_pdf_path)
556
+ output_filename = generate_filename("processed_pdf", "md")
557
+ with open(output_filename, "w", encoding="utf-8") as f:
558
+ f.write(combined_text)
559
+ st.success(f"PDF processing complete. MD file saved as {output_filename}")
560
+ st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed PDF MD"), unsafe_allow_html=True)
561
+
562
+ # ----------------- TAB: Image Process -----------------
563
+ with tab_image_process:
564
+ st.header("Image Process")
565
+ st.subheader("Upload Images for GPT-based OCR")
566
+ gpt_models = ["gpt-4o", "gpt-4o-mini"]
567
+ selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="img_gpt_model")
568
+ detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="img_detail_level")
569
+ prompt_img = st.text_input("Enter prompt for image processing", "Extract the electronic text from image", key="img_process_prompt")
570
+ uploaded_images = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader")
571
+ if st.button("Process Uploaded Images", key="process_images"):
572
+ combined_text = ""
573
+ for img_file in uploaded_images:
574
+ try:
575
+ img = Image.open(img_file)
576
+ st.image(img, caption=img_file.name)
577
+ gpt_text = process_image_with_prompt(img, prompt_img, model=selected_gpt_model, detail=detail_level)
578
+ combined_text += f"\n## {img_file.name}\n\n{gpt_text}\n"
579
+ except Exception as e:
580
+ st.error(f"Error processing image {img_file.name}: {str(e)}")
581
+ output_filename = generate_filename("processed_image", "md")
582
+ with open(output_filename, "w", encoding="utf-8") as f:
583
+ f.write(combined_text)
584
+ st.success(f"Image processing complete. MD file saved as {output_filename}")
585
+ st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed Image MD"), unsafe_allow_html=True)
586
+
587
+ # ----------------- TAB: MD Gallery -----------------
588
+ with tab_md_gallery:
589
+ st.header("MD Gallery and GPT Processing")
590
+ gpt_models = ["gpt-4o", "gpt-4o-mini"]
591
+ selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="md_gpt_model")
592
+ md_files = sorted(glob.glob("*.md"))
593
+ if md_files:
594
+ st.subheader("Individual File Processing")
595
+ cols = st.columns(2)
596
+ for idx, md_file in enumerate(md_files):
597
+ with cols[idx % 2]:
598
+ st.write(md_file)
599
+ if st.button(f"Process {md_file}", key=f"process_md_{md_file}"):
600
+ try:
601
+ with open(md_file, "r", encoding="utf-8") as f:
602
+ content = f.read()
603
+ prompt_md = "Summarize this into markdown outline with emojis and number the topics 1..12"
604
+ result_text = process_text_with_prompt(content, prompt_md, model=selected_gpt_model)
605
+ st.markdown(result_text)
606
+ output_filename = generate_filename(f"processed_{os.path.splitext(md_file)[0]}", "md")
607
+ with open(output_filename, "w", encoding="utf-8") as f:
608
+ f.write(result_text)
609
+ st.markdown(get_download_link(output_filename, "text/markdown", f"Download {output_filename}"), unsafe_allow_html=True)
610
+ except Exception as e:
611
+ st.error(f"Error processing {md_file}: {str(e)}")
612
+ st.subheader("Batch Processing")
613
+ st.write("Select MD files to combine and process:")
614
+ selected_md = {}
615
+ for md_file in md_files:
616
+ selected_md[md_file] = st.checkbox(md_file, key=f"checkbox_md_{md_file}")
617
+ batch_prompt = st.text_input("Enter batch processing prompt", "Summarize this into markdown outline with emojis and number the topics 1..12", key="batch_prompt")
618
+ if st.button("Process Selected MD Files", key="process_batch_md"):
619
+ combined_content = ""
620
+ for md_file, selected in selected_md.items():
621
+ if selected:
622
+ try:
623
+ with open(md_file, "r", encoding="utf-8") as f:
624
+ combined_content += f"\n## {md_file}\n" + f.read() + "\n"
625
+ except Exception as e:
626
+ st.error(f"Error reading {md_file}: {str(e)}")
627
+ if combined_content:
628
+ result_text = process_text_with_prompt(combined_content, batch_prompt, model=selected_gpt_model)
629
+ st.markdown(result_text)
630
+ output_filename = generate_filename("batch_processed_md", "md")
631
+ with open(output_filename, "w", encoding="utf-8") as f:
632
+ f.write(result_text)
633
+ st.success(f"Batch processing complete. MD file saved as {output_filename}")
634
+ st.markdown(get_download_link(output_filename, "text/markdown", "Download Batch Processed MD"), unsafe_allow_html=True)
635
+ else:
636
+ st.warning("No MD files selected.")
637
+ else:
638
+ st.warning("No MD files found.")
639
+
640
+ # ----------------- FINAL SIDEBAR UPDATE -----------------
641
+ # Update the asset gallery once (using its container).
642
+ def update_gallery():
643
+ container = st.session_state['asset_gallery_container']
644
+ container.empty() # Clear previous gallery content.
645
+ all_files = get_gallery_files()
646
+ if all_files:
647
+ container.markdown("### Asset Gallery 📸📖")
648
+ cols = container.columns(2)
649
+ for idx, file in enumerate(all_files[:st.session_state['gallery_size']]):
650
+ with cols[idx % 2]:
651
+ st.session_state['unique_counter'] += 1
652
+ unique_id = st.session_state['unique_counter']
653
+ if file.endswith('.png'):
654
+ st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
655
+ else:
656
+ doc = fitz.open(file)
657
+ pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
658
+ img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
659
+ st.image(img, caption=os.path.basename(file), use_container_width=True)
660
+ doc.close()
661
+ checkbox_key = f"asset_{file}_{unique_id}"
662
+ st.session_state['asset_checkboxes'][file] = st.checkbox("Use for SFT/Input", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
663
+ mime_type = "image/png" if file.endswith('.png') else "application/pdf"
664
+ st.markdown(get_download_link(file, mime_type, "Snag It! 📥"), unsafe_allow_html=True)
665
+ if st.button("Zap It! 🗑️", key=f"delete_{file}_{unique_id}"):
666
+ os.remove(file)
667
+ st.session_state['asset_checkboxes'].pop(file, None)
668
+ st.success(f"Asset {os.path.basename(file)} vaporized! 💨")
669
+ st.experimental_rerun()
670
+
671
+ # Call the gallery update once after all tabs have been processed.
672
+ update_gallery()
673
+
674
+ # Finally, update the Action Logs and History in the sidebar.
675
+ st.sidebar.subheader("Action Logs 📜")
676
+ for record in log_records:
677
+ st.sidebar.write(f"{record.asctime} - {record.levelname} - {record.message}")
678
+
679
+ st.sidebar.subheader("History 📜")
680
+ for entry in st.session_state.get("history", []):
681
+ if entry is not None:
682
+ st.sidebar.write(entry)