| | |
| | import os |
| | import glob |
| | import base64 |
| | import time |
| | import pandas as pd |
| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel |
| | from diffusers import StableDiffusionPipeline |
| | from torch.utils.data import Dataset, DataLoader |
| | import csv |
| | import fitz |
| | import requests |
| | from PIL import Image |
| | import numpy as np |
| | import logging |
| | import asyncio |
| | import aiofiles |
| | from io import BytesIO |
| | from dataclasses import dataclass |
| | from typing import Optional, Tuple |
| | import zipfile |
| | import math |
| | import random |
| | import re |
| | import gradio as gr |
| |
|
| | logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") |
| | logger = logging.getLogger(__name__) |
| | log_records = [] |
| |
|
| | class LogCaptureHandler(logging.Handler): |
| | def emit(self, record): |
| | log_records.append(record) |
| |
|
| | logger.addHandler(LogCaptureHandler()) |
| |
|
| | @dataclass |
| | class ModelConfig: |
| | name: str |
| | base_model: str |
| | size: str |
| | domain: Optional[str] = None |
| | model_type: str = "causal_lm" |
| | @property |
| | def model_path(self): |
| | return f"models/{self.name}" |
| |
|
| | @dataclass |
| | class DiffusionConfig: |
| | name: str |
| | base_model: str |
| | size: str |
| | domain: Optional[str] = None |
| | @property |
| | def model_path(self): |
| | return f"diffusion_models/{self.name}" |
| |
|
| | class ModelBuilder: |
| | def __init__(self): |
| | self.config = None |
| | self.model = None |
| | self.tokenizer = None |
| | self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! 😂", "Training complete! Time for a binary coffee break. ☕"] |
| | def load_model(self, model_path: str, config: Optional[ModelConfig] = None): |
| | self.model = AutoModelForCausalLM.from_pretrained(model_path) |
| | self.tokenizer = AutoTokenizer.from_pretrained(model_path) |
| | if self.tokenizer.pad_token is None: |
| | self.tokenizer.pad_token = self.tokenizer.eos_token |
| | if config: |
| | self.config = config |
| | self.model.to("cuda" if torch.cuda.is_available() else "cpu") |
| | return self |
| | def save_model(self, path: str): |
| | os.makedirs(os.path.dirname(path), exist_ok=True) |
| | self.model.save_pretrained(path) |
| | self.tokenizer.save_pretrained(path) |
| |
|
| | class DiffusionBuilder: |
| | def __init__(self): |
| | self.config = None |
| | self.pipeline = None |
| | def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None): |
| | self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu") |
| | if config: |
| | self.config = config |
| | return self |
| | def save_model(self, path: str): |
| | os.makedirs(os.path.dirname(path), exist_ok=True) |
| | self.pipeline.save_pretrained(path) |
| | def generate(self, prompt: str): |
| | return self.pipeline(prompt, num_inference_steps=20).images[0] |
| |
|
| | def generate_filename(sequence, ext="png"): |
| | timestamp = time.strftime("%d%m%Y%H%M%S") |
| | return f"{sequence}_{timestamp}.{ext}" |
| |
|
| | def pdf_url_to_filename(url): |
| | safe_name = re.sub(r'[<>:"/\\|?*]', '_', url) |
| | return f"{safe_name}.pdf" |
| |
|
| | def get_gallery_files(file_types=["png", "pdf"]): |
| | return sorted(list(set([f for ext in file_types for f in glob.glob(f"*.{ext}")]))) |
| |
|
| | def get_model_files(model_type="causal_lm"): |
| | path = "models/*" if model_type == "causal_lm" else "diffusion_models/*" |
| | dirs = [d for d in glob.glob(path) if os.path.isdir(d)] |
| | return dirs if dirs else ["None"] |
| |
|
| | def download_pdf(url, output_path): |
| | try: |
| | response = requests.get(url, stream=True, timeout=10) |
| | if response.status_code == 200: |
| | with open(output_path, "wb") as f: |
| | for chunk in response.iter_content(chunk_size=8192): |
| | f.write(chunk) |
| | return True |
| | except requests.RequestException as e: |
| | logger.error(f"Failed to download {url}: {e}") |
| | return False |
| |
|
| | async def process_pdf_snapshot(pdf_path, mode="single"): |
| | doc = fitz.open(pdf_path) |
| | output_files = [] |
| | if mode == "single": |
| | page = doc[0] |
| | pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | output_file = generate_filename("single", "png") |
| | pix.save(output_file) |
| | output_files.append(output_file) |
| | elif mode == "twopage": |
| | for i in range(min(2, len(doc))): |
| | page = doc[i] |
| | pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | output_file = generate_filename(f"twopage_{i}", "png") |
| | pix.save(output_file) |
| | output_files.append(output_file) |
| | elif mode == "allpages": |
| | for i in range(len(doc)): |
| | page = doc[i] |
| | pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | output_file = generate_filename(f"page_{i}", "png") |
| | pix.save(output_file) |
| | output_files.append(output_file) |
| | doc.close() |
| | return output_files |
| |
|
| | async def process_ocr(image, output_file): |
| | tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True) |
| | model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval() |
| | temp_file = f"temp_{int(time.time())}.png" |
| | image.save(temp_file) |
| | result = model.chat(tokenizer, temp_file, ocr_type='ocr') |
| | os.remove(temp_file) |
| | async with aiofiles.open(output_file, "w") as f: |
| | await f.write(result) |
| | return result |
| |
|
| | async def process_image_gen(prompt, output_file, builder): |
| | if builder and isinstance(builder, DiffusionBuilder) and builder.pipeline: |
| | pipeline = builder.pipeline |
| | else: |
| | pipeline = StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu") |
| | gen_image = pipeline(prompt, num_inference_steps=20).images[0] |
| | gen_image.save(output_file) |
| | return gen_image |
| |
|
| | |
| | def update_gallery(history, asset_checkboxes): |
| | all_files = get_gallery_files() |
| | gallery_images = [] |
| | for file in all_files[:5]: |
| | if file.endswith('.png'): |
| | gallery_images.append(Image.open(file)) |
| | else: |
| | doc = fitz.open(file) |
| | pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) |
| | img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
| | gallery_images.append(img) |
| | doc.close() |
| | history.append(f"Gallery updated: {len(all_files)} files") |
| | return gallery_images, history, asset_checkboxes |
| |
|
| | def camera_snap(image, cam_id, history, asset_checkboxes, cam_files): |
| | if image is not None: |
| | filename = generate_filename(f"cam{cam_id}") |
| | image.save(filename) |
| | history.append(f"Snapshot from Cam {cam_id}: {filename}") |
| | asset_checkboxes[filename] = True |
| | cam_files[cam_id] = filename |
| | return f"Image saved as {filename}", Image.open(filename), history, asset_checkboxes, cam_files |
| | elif cam_files.get(cam_id) and os.path.exists(cam_files[cam_id]): |
| | return f"Showing previous capture: {cam_files[cam_id]}", Image.open(cam_files[cam_id]), history, asset_checkboxes, cam_files |
| | return "No image captured", None, history, asset_checkboxes, cam_files |
| |
|
| | def download_pdfs(urls, history, asset_checkboxes): |
| | urls = urls.strip().split("\n") |
| | downloaded = [] |
| | for url in urls: |
| | if url: |
| | output_path = pdf_url_to_filename(url) |
| | if download_pdf(url, output_path): |
| | downloaded.append(output_path) |
| | history.append(f"Downloaded PDF: {output_path}") |
| | asset_checkboxes[output_path] = True |
| | return f"Downloaded {len(downloaded)} PDFs", history, asset_checkboxes |
| |
|
| | def upload_pdfs(pdf_files, history, asset_checkboxes): |
| | uploaded = [] |
| | for pdf_file in pdf_files: |
| | if pdf_file: |
| | output_path = f"uploaded_{int(time.time())}_{pdf_file.name}" |
| | with open(output_path, "wb") as f: |
| | f.write(pdf_file.read()) |
| | uploaded.append(output_path) |
| | history.append(f"Uploaded PDF: {output_path}") |
| | asset_checkboxes[output_path] = True |
| | return f"Uploaded {len(uploaded)} PDFs", history, asset_checkboxes |
| |
|
| | def snapshot_pdfs(mode, history, asset_checkboxes): |
| | selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and asset_checkboxes.get(path, False)] |
| | if not selected_pdfs: |
| | return "No PDFs selected", [], history, asset_checkboxes |
| | snapshots = [] |
| | mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages"}[mode] |
| | for pdf_path in selected_pdfs: |
| | snap_files = asyncio.run(process_pdf_snapshot(pdf_path, mode_key)) |
| | for snap in snap_files: |
| | snapshots.append(Image.open(snap)) |
| | asset_checkboxes[snap] = True |
| | history.append(f"Snapshot {mode_key}: {snap}") |
| | return f"Generated {len(snapshots)} snapshots", snapshots, history, asset_checkboxes |
| |
|
| | def process_ocr_all(history, asset_checkboxes): |
| | all_files = get_gallery_files() |
| | if not all_files: |
| | return "No assets to OCR", history, asset_checkboxes |
| | full_text = "# OCR Results\n\n" |
| | for file in all_files: |
| | if file.endswith('.png'): |
| | image = Image.open(file) |
| | else: |
| | doc = fitz.open(file) |
| | pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
| | doc.close() |
| | output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt") |
| | result = asyncio.run(process_ocr(image, output_file)) |
| | full_text += f"## {os.path.basename(file)}\n\n{result}\n\n" |
| | history.append(f"OCR Test: {file} -> {output_file}") |
| | md_output_file = f"full_ocr_{int(time.time())}.md" |
| | with open(md_output_file, "w") as f: |
| | f.write(full_text) |
| | return f"Full OCR saved to {md_output_file}", history, asset_checkboxes |
| |
|
| | def process_ocr_single(file_path, history, asset_checkboxes): |
| | if not file_path: |
| | return "No file selected", None, "", history, asset_checkboxes |
| | if file_path.endswith('.png'): |
| | image = Image.open(file_path) |
| | else: |
| | doc = fitz.open(file_path) |
| | pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
| | doc.close() |
| | output_file = generate_filename("ocr_output", "txt") |
| | result = asyncio.run(process_ocr(image, output_file)) |
| | history.append(f"OCR Test: {file_path} -> {output_file}") |
| | return f"OCR output saved to {output_file}", image, result, history, asset_checkboxes |
| |
|
| | def build_model(model_type, base_model, model_name, domain, history): |
| | config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model, size="small", domain=domain) |
| | builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder() |
| | builder.load_model(base_model, config) |
| | builder.save_model(config.model_path) |
| | history.append(f"Built {model_type} model: {model_name}") |
| | return builder, f"Model saved to {config.model_path}", history |
| |
|
| | def image_gen(prompt, file_path, builder, history, asset_checkboxes): |
| | if not file_path: |
| | return "No file selected", None, history, asset_checkboxes |
| | if file_path.endswith('.png'): |
| | image = Image.open(file_path) |
| | else: |
| | doc = fitz.open(file_path) |
| | pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
| | image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
| | doc.close() |
| | output_file = generate_filename("gen_output", "png") |
| | gen_image = asyncio.run(process_image_gen(prompt, output_file, builder)) |
| | history.append(f"Image Gen Test: {prompt} -> {output_file}") |
| | asset_checkboxes[output_file] = True |
| | return f"Image saved to {output_file}", gen_image, history, asset_checkboxes |
| |
|
| | |
| | with gr.Blocks(title="AI Vision & SFT Titans 🚀") as demo: |
| | gr.Markdown("# AI Vision & SFT Titans 🚀") |
| | history = gr.State(value=[]) |
| | builder = gr.State(value=None) |
| | asset_checkboxes = gr.State(value={}) |
| | cam_files = gr.State(value={}) |
| |
|
| | with gr.Row(): |
| | with gr.Column(scale=1): |
| | gr.Markdown("## Captured Files 📜") |
| | gallery_output = gr.Gallery(label="Asset Gallery", columns=2, height="auto") |
| | gr.Button("Update Gallery").click(update_gallery, inputs=[history, asset_checkboxes], outputs=[gallery_output, history, asset_checkboxes]) |
| | gr.Markdown("## History 📜") |
| | history_output = gr.Textbox(label="History", lines=5, interactive=False) |
| | gr.Markdown("## Action Logs 📜") |
| | log_output = gr.Textbox(label="Logs", value="\n".join([f"{r.asctime} - {r.levelname} - {r.message}" for r in log_records]), lines=5, interactive=False) |
| |
|
| | with gr.Column(scale=3): |
| | with gr.Tabs(): |
| | with gr.TabItem("Camera Snap 📷"): |
| | with gr.Row(): |
| | cam0_input = gr.Image(type="pil", label="Camera 0") |
| | cam1_input = gr.Image(type="pil", label="Camera 1") |
| | with gr.Row(): |
| | cam0_output = gr.Textbox(label="Cam 0 Status") |
| | cam1_output = gr.Textbox(label="Cam 1 Status") |
| | with gr.Row(): |
| | cam0_image = gr.Image(label="Cam 0 Preview") |
| | cam1_image = gr.Image(label="Cam 1 Preview") |
| | gr.Button("Capture Cam 0").click(camera_snap, inputs=[cam0_input, gr.State(value=0), history, asset_checkboxes, cam_files], outputs=[cam0_output, cam0_image, history, asset_checkboxes, cam_files]) |
| | gr.Button("Capture Cam 1").click(camera_snap, inputs=[cam1_input, gr.State(value=1), history, asset_checkboxes, cam_files], outputs=[cam1_output, cam1_image, history, asset_checkboxes, cam_files]) |
| |
|
| | with gr.TabItem("Download PDFs 📥"): |
| | url_input = gr.Textbox(label="Enter PDF URLs (one per line)", lines=5) |
| | pdf_upload = gr.File(label="Upload PDFs", file_count="multiple", type="binary") |
| | pdf_output = gr.Textbox(label="Status") |
| | snapshot_mode = gr.Dropdown(["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], label="Snapshot Mode") |
| | snapshot_output = gr.Textbox(label="Snapshot Status") |
| | snapshot_images = gr.Gallery(label="Snapshots", columns=2, height="auto") |
| | gr.Button("Download URLs").click(download_pdfs, inputs=[url_input, history, asset_checkboxes], outputs=[pdf_output, history, asset_checkboxes]) |
| | gr.Button("Upload PDFs").click(upload_pdfs, inputs=[pdf_upload, history, asset_checkboxes], outputs=[pdf_output, history, asset_checkboxes]) |
| | gr.Button("Snapshot Selected").click(snapshot_pdfs, inputs=[snapshot_mode, history, asset_checkboxes], outputs=[snapshot_output, snapshot_images, history, asset_checkboxes]) |
| |
|
| | with gr.TabItem("Test OCR 🔍"): |
| | all_files = gr.Dropdown(choices=get_gallery_files(), label="Select File") |
| | ocr_output = gr.Textbox(label="Status") |
| | ocr_image = gr.Image(label="Input Image") |
| | ocr_result = gr.Textbox(label="OCR Result", lines=5) |
| | gr.Button("OCR All Assets").click(process_ocr_all, inputs=[history, asset_checkboxes], outputs=[ocr_output, history, asset_checkboxes]) |
| | gr.Button("OCR Selected").click(process_ocr_single, inputs=[all_files, history, asset_checkboxes], outputs=[ocr_output, ocr_image, ocr_result, history, asset_checkboxes]) |
| |
|
| | with gr.TabItem("Build Titan 🌱"): |
| | model_type = gr.Dropdown(["Causal LM", "Diffusion"], label="Model Type") |
| | base_model = gr.Dropdown( |
| | choices=["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"], |
| | label="Base Model", |
| | value="HuggingFaceTB/SmolLM-135M" |
| | ) |
| | model_name = gr.Textbox(label="Model Name", value=f"tiny-titan-{int(time.time())}") |
| | domain = gr.Textbox(label="Target Domain", value="general") |
| | build_output = gr.Textbox(label="Status") |
| | gr.Button("Build").click(build_model, inputs=[model_type, base_model, model_name, domain, history], outputs=[builder, build_output, history]) |
| |
|
| | with gr.TabItem("Test Image Gen 🎨"): |
| | gen_file = gr.Dropdown(choices=get_gallery_files(), label="Select Reference File") |
| | gen_prompt = gr.Textbox(label="Prompt", value="Generate a neon superhero version of this image") |
| | gen_output = gr.Textbox(label="Status") |
| | gen_image = gr.Image(label="Generated Image") |
| | gr.Button("Generate").click(image_gen, inputs=[gen_prompt, gen_file, builder, history, asset_checkboxes], outputs=[gen_output, gen_image, history, asset_checkboxes]) |
| |
|
| | |
| | demo.load(lambda h: "\n".join(h[-5:]), inputs=[history], outputs=[history_output]) |
| |
|
| | demo.launch() |
