|
|
|
import os |
|
import glob |
|
import base64 |
|
import time |
|
import shutil |
|
import streamlit as st |
|
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 cv2 |
|
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 |
|
|
|
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()) |
|
|
|
st.set_page_config( |
|
page_title="AI Vision & SFT Titans 🚀", |
|
page_icon="🤖", |
|
layout="wide", |
|
initial_sidebar_state="expanded", |
|
menu_items={ |
|
'Get Help': 'https://huggingface.co/awacke1', |
|
'Report a Bug': 'https://huggingface.co/spaces/awacke1', |
|
'About': "AI Vision & SFT Titans: PDFs, OCR, Image Gen, Line Drawings, Custom Diffusion, and SFT on CPU! 🌌" |
|
} |
|
) |
|
|
|
if 'history' not in st.session_state: |
|
st.session_state['history'] = [] |
|
if 'builder' not in st.session_state: |
|
st.session_state['builder'] = None |
|
if 'model_loaded' not in st.session_state: |
|
st.session_state['model_loaded'] = False |
|
if 'processing' not in st.session_state: |
|
st.session_state['processing'] = {} |
|
if 'pdf_checkboxes' not in st.session_state: |
|
st.session_state['pdf_checkboxes'] = {} |
|
if 'downloaded_pdfs' not in st.session_state: |
|
st.session_state['downloaded_pdfs'] = {} |
|
|
|
@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 SFTDataset(Dataset): |
|
def __init__(self, data, tokenizer, max_length=128): |
|
self.data = data |
|
self.tokenizer = tokenizer |
|
self.max_length = max_length |
|
def __len__(self): |
|
return len(self.data) |
|
def __getitem__(self, idx): |
|
prompt = self.data[idx]["prompt"] |
|
response = self.data[idx]["response"] |
|
full_text = f"{prompt} {response}" |
|
full_encoding = self.tokenizer(full_text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt") |
|
prompt_encoding = self.tokenizer(prompt, max_length=self.max_length, padding=False, truncation=True, return_tensors="pt") |
|
input_ids = full_encoding["input_ids"].squeeze() |
|
attention_mask = full_encoding["attention_mask"].squeeze() |
|
labels = input_ids.clone() |
|
prompt_len = prompt_encoding["input_ids"].shape[1] |
|
if prompt_len < self.max_length: |
|
labels[:prompt_len] = -100 |
|
return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels} |
|
|
|
class DiffusionDataset(Dataset): |
|
def __init__(self, images, texts): |
|
self.images = images |
|
self.texts = texts |
|
def __len__(self): |
|
return len(self.images) |
|
def __getitem__(self, idx): |
|
return {"image": self.images[idx], "text": self.texts[idx]} |
|
|
|
class TinyDiffusionDataset(Dataset): |
|
def __init__(self, images): |
|
self.images = [torch.tensor(np.array(img.convert("RGB")).transpose(2, 0, 1), dtype=torch.float32) / 255.0 for img in images] |
|
def __len__(self): |
|
return len(self.images) |
|
def __getitem__(self, idx): |
|
return self.images[idx] |
|
|
|
class TinyUNet(nn.Module): |
|
def __init__(self, in_channels=3, out_channels=3): |
|
super(TinyUNet, self).__init__() |
|
self.down1 = nn.Conv2d(in_channels, 32, 3, padding=1) |
|
self.down2 = nn.Conv2d(32, 64, 3, padding=1, stride=2) |
|
self.mid = nn.Conv2d(64, 128, 3, padding=1) |
|
self.up1 = nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1) |
|
self.up2 = nn.Conv2d(64 + 32, 32, 3, padding=1) |
|
self.out = nn.Conv2d(32, out_channels, 3, padding=1) |
|
self.time_embed = nn.Linear(1, 64) |
|
|
|
def forward(self, x, t): |
|
t_embed = F.relu(self.time_embed(t.unsqueeze(-1))) |
|
t_embed = t_embed.view(t_embed.size(0), t_embed.size(1), 1, 1) |
|
|
|
x1 = F.relu(self.down1(x)) |
|
x2 = F.relu(self.down2(x1)) |
|
x_mid = F.relu(self.mid(x2)) + t_embed |
|
x_up1 = F.relu(self.up1(x_mid)) |
|
x_up2 = F.relu(self.up2(torch.cat([x_up1, x1], dim=1))) |
|
return self.out(x_up2) |
|
|
|
class TinyDiffusion: |
|
def __init__(self, model, timesteps=100): |
|
self.model = model |
|
self.timesteps = timesteps |
|
self.beta = torch.linspace(0.0001, 0.02, timesteps) |
|
self.alpha = 1 - self.beta |
|
self.alpha_cumprod = torch.cumprod(self.alpha, dim=0) |
|
|
|
def train(self, images, epochs=50): |
|
dataset = TinyDiffusionDataset(images) |
|
dataloader = DataLoader(dataset, batch_size=1, shuffle=True) |
|
optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4) |
|
device = torch.device("cpu") |
|
self.model.to(device) |
|
for epoch in range(epochs): |
|
total_loss = 0 |
|
for x in dataloader: |
|
x = x.to(device) |
|
t = torch.randint(0, self.timesteps, (x.size(0),), device=device).float() |
|
noise = torch.randn_like(x) |
|
alpha_t = self.alpha_cumprod[t.long()].view(-1, 1, 1, 1) |
|
x_noisy = torch.sqrt(alpha_t) * x + torch.sqrt(1 - alpha_t) * noise |
|
pred_noise = self.model(x_noisy, t) |
|
loss = F.mse_loss(pred_noise, noise) |
|
optimizer.zero_grad() |
|
loss.backward() |
|
optimizer.step() |
|
total_loss += loss.item() |
|
logger.info(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(dataloader):.4f}") |
|
return self |
|
|
|
def generate(self, size=(64, 64), steps=100): |
|
device = torch.device("cpu") |
|
x = torch.randn(1, 3, size[0], size[1], device=device) |
|
for t in reversed(range(steps)): |
|
t_tensor = torch.full((1,), t, device=device, dtype=torch.float32) |
|
alpha_t = self.alpha_cumprod[t].view(-1, 1, 1, 1) |
|
pred_noise = self.model(x, t_tensor) |
|
x = (x - (1 - self.alpha[t]) / torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(self.alpha[t]) |
|
if t > 0: |
|
x += torch.sqrt(self.beta[t]) * torch.randn_like(x) |
|
x = torch.clamp(x * 255, 0, 255).byte() |
|
return Image.fromarray(x.squeeze(0).permute(1, 2, 0).cpu().numpy()) |
|
|
|
def upscale(self, image, scale_factor=2): |
|
img_tensor = torch.tensor(np.array(image.convert("RGB")).transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0) / 255.0 |
|
upscaled = F.interpolate(img_tensor, scale_factor=scale_factor, mode='bilinear', align_corners=False) |
|
upscaled = torch.clamp(upscaled * 255, 0, 255).byte() |
|
return Image.fromarray(upscaled.squeeze(0).permute(1, 2, 0).cpu().numpy()) |
|
|
|
class ModelBuilder: |
|
def __init__(self): |
|
self.config = None |
|
self.model = None |
|
self.tokenizer = None |
|
self.sft_data = 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): |
|
with st.spinner(f"Loading {model_path}... ⏳"): |
|
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") |
|
st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}") |
|
return self |
|
def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4): |
|
self.sft_data = [] |
|
with open(csv_path, "r") as f: |
|
reader = csv.DictReader(f) |
|
for row in reader: |
|
self.sft_data.append({"prompt": row["prompt"], "response": row["response"]}) |
|
dataset = SFTDataset(self.sft_data, self.tokenizer) |
|
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) |
|
optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5) |
|
self.model.train() |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
self.model.to(device) |
|
for epoch in range(epochs): |
|
with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️"): |
|
total_loss = 0 |
|
for batch in dataloader: |
|
optimizer.zero_grad() |
|
input_ids = batch["input_ids"].to(device) |
|
attention_mask = batch["attention_mask"].to(device) |
|
labels = batch["labels"].to(device) |
|
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) |
|
loss = outputs.loss |
|
loss.backward() |
|
optimizer.step() |
|
total_loss += loss.item() |
|
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}") |
|
st.success(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}") |
|
return self |
|
def save_model(self, path: str): |
|
with st.spinner("Saving model... 💾"): |
|
os.makedirs(os.path.dirname(path), exist_ok=True) |
|
self.model.save_pretrained(path) |
|
self.tokenizer.save_pretrained(path) |
|
st.success(f"Model saved at {path}! ✅") |
|
def evaluate(self, prompt: str, status_container=None): |
|
self.model.eval() |
|
if status_container: |
|
status_container.write("Preparing to evaluate... 🧠") |
|
try: |
|
with torch.no_grad(): |
|
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device) |
|
outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7) |
|
return self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
except Exception as e: |
|
if status_container: |
|
status_container.error(f"Oops! Something broke: {str(e)} 💥") |
|
return f"Error: {str(e)}" |
|
|
|
class DiffusionBuilder: |
|
def __init__(self): |
|
self.config = None |
|
self.pipeline = None |
|
def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None): |
|
with st.spinner(f"Loading diffusion model {model_path}... ⏳"): |
|
self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu") |
|
if config: |
|
self.config = config |
|
st.success(f"Diffusion model loaded! 🎨") |
|
return self |
|
def fine_tune_sft(self, images, texts, epochs=3): |
|
dataset = DiffusionDataset(images, texts) |
|
dataloader = DataLoader(dataset, batch_size=1, shuffle=True) |
|
optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5) |
|
self.pipeline.unet.train() |
|
for epoch in range(epochs): |
|
with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"): |
|
total_loss = 0 |
|
for batch in dataloader: |
|
optimizer.zero_grad() |
|
image = batch["image"][0].to(self.pipeline.device) |
|
text = batch["text"][0] |
|
latents = self.pipeline.vae.encode(torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device)).latent_dist.sample() |
|
noise = torch.randn_like(latents) |
|
timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device) |
|
noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps) |
|
text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0] |
|
pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample |
|
loss = torch.nn.functional.mse_loss(pred_noise, noise) |
|
loss.backward() |
|
optimizer.step() |
|
total_loss += loss.item() |
|
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}") |
|
st.success("Diffusion SFT Fine-tuning completed! 🎨") |
|
return self |
|
def save_model(self, path: str): |
|
with st.spinner("Saving diffusion model... 💾"): |
|
os.makedirs(os.path.dirname(path), exist_ok=True) |
|
self.pipeline.save_pretrained(path) |
|
st.success(f"Diffusion model saved at {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_download_link(file_path, mime_type="application/pdf", label="Download"): |
|
with open(file_path, 'rb') as f: |
|
data = f.read() |
|
b64 = base64.b64encode(data).decode() |
|
return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>' |
|
|
|
def zip_directory(directory_path, zip_path): |
|
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: |
|
for root, _, files in os.walk(directory_path): |
|
for file in files: |
|
zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path))) |
|
|
|
def get_model_files(model_type="causal_lm"): |
|
path = "models/*" if model_type == "causal_lm" else "diffusion_models/*" |
|
return [d for d in glob.glob(path) if os.path.isdir(d)] |
|
|
|
def get_gallery_files(file_types=["png"]): |
|
return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")]) |
|
|
|
def get_pdf_files(): |
|
return sorted(glob.glob("*.pdf")) |
|
|
|
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"): |
|
start_time = time.time() |
|
status = st.empty() |
|
status.text(f"Processing PDF Snapshot ({mode})... (0s)") |
|
try: |
|
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 == "allthumbs": |
|
for i in range(len(doc)): |
|
page = doc[i] |
|
pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) |
|
output_file = generate_filename(f"thumb_{i}", "png") |
|
pix.save(output_file) |
|
output_files.append(output_file) |
|
doc.close() |
|
elapsed = int(time.time() - start_time) |
|
status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!") |
|
update_gallery() |
|
return output_files |
|
except Exception as e: |
|
status.error(f"Failed to process PDF: {str(e)}") |
|
return [] |
|
|
|
async def process_ocr(image, output_file): |
|
start_time = time.time() |
|
status = st.empty() |
|
status.text("Processing GOT-OCR2_0... (0s)") |
|
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() |
|
result = model.chat(tokenizer, image, ocr_type='ocr') |
|
elapsed = int(time.time() - start_time) |
|
status.text(f"GOT-OCR2_0 completed in {elapsed}s!") |
|
async with aiofiles.open(output_file, "w") as f: |
|
await f.write(result) |
|
update_gallery() |
|
return result |
|
|
|
async def process_image_gen(prompt, output_file): |
|
start_time = time.time() |
|
status = st.empty() |
|
status.text("Processing Image Gen... (0s)") |
|
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] |
|
elapsed = int(time.time() - start_time) |
|
status.text(f"Image Gen completed in {elapsed}s!") |
|
gen_image.save(output_file) |
|
update_gallery() |
|
return gen_image |
|
|
|
async def process_custom_diffusion(images, output_file, model_name): |
|
start_time = time.time() |
|
status = st.empty() |
|
status.text(f"Training {model_name}... (0s)") |
|
unet = TinyUNet() |
|
diffusion = TinyDiffusion(unet) |
|
diffusion.train(images) |
|
gen_image = diffusion.generate() |
|
upscaled_image = diffusion.upscale(gen_image, scale_factor=2) |
|
elapsed = int(time.time() - start_time) |
|
status.text(f"{model_name} completed in {elapsed}s!") |
|
upscaled_image.save(output_file) |
|
update_gallery() |
|
return upscaled_image |
|
|
|
def mock_search(query: str) -> str: |
|
if "superhero" in query.lower(): |
|
return "Latest trends: Gold-plated Batman statues, VR superhero battles." |
|
return "No relevant results found." |
|
|
|
def mock_duckduckgo_search(query: str) -> str: |
|
if "superhero party trends" in query.lower(): |
|
return """ |
|
Latest trends for 2025: |
|
- Luxury decorations: Gold-plated Batman statues, holographic Avengers displays. |
|
- Entertainment: Live stunt shows with Iron Man suits, VR superhero battles. |
|
- Catering: Gourmet kryptonite-green cocktails, Thor’s hammer-shaped appetizers. |
|
""" |
|
return "No relevant results found." |
|
|
|
class PartyPlannerAgent: |
|
def __init__(self, model, tokenizer): |
|
self.model = model |
|
self.tokenizer = tokenizer |
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
self.model.to(self.device) |
|
def generate(self, prompt: str) -> str: |
|
self.model.eval() |
|
with torch.no_grad(): |
|
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device) |
|
outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7) |
|
return self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
def plan_party(self, task: str) -> pd.DataFrame: |
|
search_result = mock_duckduckgo_search("latest superhero party trends") |
|
prompt = f"Given this context: '{search_result}'\n{task}" |
|
plan_text = self.generate(prompt) |
|
locations = { |
|
"Wayne Manor": (42.3601, -71.0589), |
|
"New York": (40.7128, -74.0060), |
|
"Los Angeles": (34.0522, -118.2437), |
|
"London": (51.5074, -0.1278) |
|
} |
|
wayne_coords = locations["Wayne Manor"] |
|
travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"} |
|
catchphrases = ["To the Batmobile!", "Avengers, assemble!", "I am Iron Man!", "By the power of Grayskull!"] |
|
data = [ |
|
{"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues", "Catchphrase": random.choice(catchphrases)}, |
|
{"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Holographic Avengers displays", "Catchphrase": random.choice(catchphrases)}, |
|
{"Location": "London", "Travel Time (hrs)": travel_times["London"], "Luxury Idea": "Live stunt shows with Iron Man suits", "Catchphrase": random.choice(catchphrases)}, |
|
{"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles", "Catchphrase": random.choice(catchphrases)}, |
|
{"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gourmet kryptonite-green cocktails", "Catchphrase": random.choice(catchphrases)}, |
|
{"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Thor’s hammer-shaped appetizers", "Catchphrase": random.choice(catchphrases)}, |
|
] |
|
return pd.DataFrame(data) |
|
|
|
class CVPartyPlannerAgent: |
|
def __init__(self, pipeline): |
|
self.pipeline = pipeline |
|
def generate(self, prompt: str) -> Image.Image: |
|
return self.pipeline(prompt, num_inference_steps=20).images[0] |
|
def plan_party(self, task: str) -> pd.DataFrame: |
|
search_result = mock_search("superhero party trends") |
|
prompt = f"Given this context: '{search_result}'\n{task}" |
|
data = [ |
|
{"Theme": "Batman", "Image Idea": "Gold-plated Batman statue"}, |
|
{"Theme": "Avengers", "Image Idea": "VR superhero battle scene"} |
|
] |
|
return pd.DataFrame(data) |
|
|
|
def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_coords: Tuple[float, float], cruising_speed_kmh: float = 750.0) -> float: |
|
def to_radians(degrees: float) -> float: |
|
return degrees * (math.pi / 180) |
|
lat1, lon1 = map(to_radians, origin_coords) |
|
lat2, lon2 = map(to_radians, destination_coords) |
|
EARTH_RADIUS_KM = 6371.0 |
|
dlon = lon2 - lon1 |
|
dlat = lat2 - lat1 |
|
a = (math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2) |
|
c = 2 * math.asin(math.sqrt(a)) |
|
distance = EARTH_RADIUS_KM * c |
|
actual_distance = distance * 1.1 |
|
flight_time = (actual_distance / cruising_speed_kmh) + 1.0 |
|
return round(flight_time, 2) |
|
|
|
st.title("AI Vision & SFT Titans 🚀") |
|
|
|
st.sidebar.header("Captured Files 📜") |
|
gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2) |
|
def update_gallery(): |
|
media_files = get_gallery_files(["png"]) |
|
pdf_files = get_pdf_files() |
|
if media_files or pdf_files: |
|
st.sidebar.subheader("Images 📸") |
|
cols = st.sidebar.columns(2) |
|
for idx, file in enumerate(media_files[:gallery_size * 2]): |
|
with cols[idx % 2]: |
|
st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True) |
|
st.sidebar.subheader("PDF Downloads 📖") |
|
for pdf_file in pdf_files[:gallery_size * 2]: |
|
st.markdown(get_download_link(pdf_file, "application/pdf", f"📥 Grab {os.path.basename(pdf_file)}"), unsafe_allow_html=True) |
|
update_gallery() |
|
|
|
st.sidebar.subheader("Model Management 🗂️") |
|
model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"], key="sidebar_model_type") |
|
model_dirs = get_model_files(model_type) |
|
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs, key="sidebar_model_select") |
|
if selected_model != "None" and st.sidebar.button("Load Model 📂"): |
|
builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder() |
|
config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small") |
|
builder.load_model(selected_model, config) |
|
st.session_state['builder'] = builder |
|
st.session_state['model_loaded'] = True |
|
st.rerun() |
|
|
|
st.sidebar.subheader("Action Logs 📜") |
|
log_container = st.sidebar.empty() |
|
with log_container: |
|
for record in log_records: |
|
st.write(f"{record.asctime} - {record.levelname} - {record.message}") |
|
|
|
st.sidebar.subheader("History 📜") |
|
history_container = st.sidebar.empty() |
|
with history_container: |
|
for entry in st.session_state['history'][-gallery_size * 2:]: |
|
st.write(entry) |
|
|
|
tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8, tab9 = st.tabs([ |
|
"Camera Snap 📷", "Download PDFs 📥", "Build Titan 🌱", "Fine-Tune Titan 🔧", |
|
"Test Titan 🧪", "Agentic RAG Party 🌐", "Test OCR 🔍", "Test Image Gen 🎨", "Custom Diffusion 🎨🤓" |
|
]) |
|
|
|
with tab1: |
|
st.header("Camera Snap 📷") |
|
st.subheader("Single Capture") |
|
cols = st.columns(2) |
|
with cols[0]: |
|
cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0") |
|
if cam0_img: |
|
filename = generate_filename("cam0") |
|
with open(filename, "wb") as f: |
|
f.write(cam0_img.getvalue()) |
|
entry = f"Snapshot from Cam 0: {filename}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 0:")] + [entry] |
|
st.image(Image.open(filename), caption="Camera 0", use_container_width=True) |
|
logger.info(f"Saved snapshot from Camera 0: {filename}") |
|
update_gallery() |
|
with cols[1]: |
|
cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1") |
|
if cam1_img: |
|
filename = generate_filename("cam1") |
|
with open(filename, "wb") as f: |
|
f.write(cam1_img.getvalue()) |
|
entry = f"Snapshot from Cam 1: {filename}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 1:")] + [entry] |
|
st.image(Image.open(filename), caption="Camera 1", use_container_width=True) |
|
logger.info(f"Saved snapshot from Camera 1: {filename}") |
|
update_gallery() |
|
|
|
with tab2: |
|
st.header("Download PDFs 📥") |
|
if st.button("Examples 📚"): |
|
example_urls = [ |
|
"https://arxiv.org/pdf/2308.03892", |
|
"https://arxiv.org/pdf/1912.01703", |
|
"https://arxiv.org/pdf/2408.11039", |
|
"https://arxiv.org/pdf/2109.10282", |
|
"https://arxiv.org/pdf/2112.10752", |
|
"https://arxiv.org/pdf/2308.11236", |
|
"https://arxiv.org/pdf/1706.03762", |
|
"https://arxiv.org/pdf/2006.11239", |
|
"https://arxiv.org/pdf/2305.11207", |
|
"https://arxiv.org/pdf/2106.09685", |
|
"https://arxiv.org/pdf/2005.11401", |
|
"https://arxiv.org/pdf/2106.10504" |
|
] |
|
st.session_state['pdf_urls'] = "\n".join(example_urls) |
|
|
|
url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200) |
|
if st.button("Robo-Download 🤖"): |
|
urls = url_input.strip().split("\n") |
|
progress_bar = st.progress(0) |
|
status_text = st.empty() |
|
total_urls = len(urls) |
|
existing_pdfs = get_pdf_files() |
|
for idx, url in enumerate(urls): |
|
if url: |
|
output_path = pdf_url_to_filename(url) |
|
status_text.text(f"Fetching {idx + 1}/{total_urls}: {os.path.basename(output_path)}...") |
|
if output_path not in existing_pdfs: |
|
if download_pdf(url, output_path): |
|
st.session_state['downloaded_pdfs'][url] = output_path |
|
logger.info(f"Downloaded PDF from {url} to {output_path}") |
|
entry = f"Downloaded PDF: {output_path}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
else: |
|
st.error(f"Failed to nab {url} 😿") |
|
else: |
|
st.info(f"Already got {os.path.basename(output_path)}! Skipping... 🐾") |
|
st.session_state['downloaded_pdfs'][url] = output_path |
|
progress_bar.progress((idx + 1) / total_urls) |
|
status_text.text("Robo-Download complete! 🚀") |
|
update_gallery() |
|
|
|
st.subheader("PDF Gallery 📖") |
|
downloaded_pdfs = list(st.session_state['downloaded_pdfs'].values()) |
|
if downloaded_pdfs: |
|
cols_per_row = 3 |
|
for i in range(0, len(downloaded_pdfs), cols_per_row): |
|
cols = st.columns(cols_per_row) |
|
for j, pdf_path in enumerate(downloaded_pdfs[i:i + cols_per_row]): |
|
with cols[j]: |
|
doc = fitz.open(pdf_path) |
|
page = doc[0] |
|
pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) |
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
|
st.image(img, caption=os.path.basename(pdf_path), use_container_width=True) |
|
checkbox_key = f"pdf_{pdf_path}" |
|
st.session_state['pdf_checkboxes'][checkbox_key] = st.checkbox( |
|
"Use for SFT/Input", |
|
value=st.session_state['pdf_checkboxes'].get(checkbox_key, False), |
|
key=checkbox_key |
|
) |
|
st.markdown(get_download_link(pdf_path, "application/pdf", "Snag It! 📥"), unsafe_allow_html=True) |
|
if st.button("Zap It! 🗑️", key=f"delete_{pdf_path}"): |
|
os.remove(pdf_path) |
|
url_key = next((k for k, v in st.session_state['downloaded_pdfs'].items() if v == pdf_path), None) |
|
if url_key: |
|
del st.session_state['downloaded_pdfs'][url_key] |
|
del st.session_state['pdf_checkboxes'][checkbox_key] |
|
st.success(f"PDF {os.path.basename(pdf_path)} vaporized! 💨") |
|
st.rerun() |
|
doc.close() |
|
else: |
|
st.info("No PDFs captured yet. Feed the robo-downloader some URLs! 🤖") |
|
|
|
mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (Thumbnails)"], key="download_mode") |
|
if st.button("Snapshot Selected 📸"): |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
if selected_pdfs: |
|
for pdf_path in selected_pdfs: |
|
mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (Thumbnails)": "allthumbs"}[mode] |
|
snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key)) |
|
for snapshot in snapshots: |
|
st.image(Image.open(snapshot), caption=snapshot, use_container_width=True) |
|
else: |
|
st.warning("No PDFs selected for snapshotting! Check some boxes first. 📝") |
|
|
|
with tab3: |
|
st.header("Build Titan 🌱") |
|
model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type") |
|
base_model = st.selectbox("Select Tiny Model", |
|
["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else |
|
["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"]) |
|
model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}") |
|
domain = st.text_input("Target Domain", "general") |
|
if st.button("Download Model ⬇️"): |
|
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) |
|
st.session_state['builder'] = builder |
|
st.session_state['model_loaded'] = True |
|
entry = f"Built {model_type} model: {model_name}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.success(f"Model downloaded and saved to {config.model_path}! 🎉") |
|
st.rerun() |
|
|
|
with tab4: |
|
st.header("Fine-Tune Titan 🔧") |
|
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False): |
|
st.warning("Please build or load a Titan first! ⚠️") |
|
else: |
|
if isinstance(st.session_state['builder'], ModelBuilder): |
|
if st.button("Generate Sample CSV 📝"): |
|
sample_data = [ |
|
{"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human smarts in machines."}, |
|
{"prompt": "Explain machine learning", "response": "Machine learning is AI’s gym where models bulk up on data."}, |
|
] |
|
csv_path = f"sft_data_{int(time.time())}.csv" |
|
with open(csv_path, "w", newline="") as f: |
|
writer = csv.DictWriter(f, fieldnames=["prompt", "response"]) |
|
writer.writeheader() |
|
writer.writerows(sample_data) |
|
st.markdown(get_download_link(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True) |
|
st.success(f"Sample CSV generated as {csv_path}! ✅") |
|
|
|
uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv") |
|
if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"): |
|
csv_path = f"uploaded_sft_data_{int(time.time())}.csv" |
|
with open(csv_path, "wb") as f: |
|
f.write(uploaded_csv.read()) |
|
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}" |
|
new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small", domain=st.session_state['builder'].config.domain) |
|
st.session_state['builder'].config = new_config |
|
st.session_state['builder'].fine_tune_sft(csv_path) |
|
st.session_state['builder'].save_model(new_config.model_path) |
|
zip_path = f"{new_config.model_path}.zip" |
|
zip_directory(new_config.model_path, zip_path) |
|
entry = f"Fine-tuned Causal LM: {new_model_name}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True) |
|
st.rerun() |
|
elif isinstance(st.session_state['builder'], DiffusionBuilder): |
|
captured_files = get_gallery_files(["png"]) |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
if len(captured_files) + len(selected_pdfs) >= 2: |
|
demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files] |
|
for pdf_path in selected_pdfs: |
|
demo_data.append({"image": pdf_path, "text": f"PDF {os.path.basename(pdf_path)}"}) |
|
edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic") |
|
if st.button("Fine-Tune with Dataset 🔄"): |
|
images = [Image.open(row["image"]) if row["image"].endswith('.png') else Image.frombytes("RGB", fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).size, fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).samples) for _, row in edited_data.iterrows()] |
|
texts = [row["text"] for _, row in edited_data.iterrows()] |
|
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}" |
|
new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small", domain=st.session_state['builder'].config.domain) |
|
st.session_state['builder'].config = new_config |
|
st.session_state['builder'].fine_tune_sft(images, texts) |
|
st.session_state['builder'].save_model(new_config.model_path) |
|
zip_path = f"{new_config.model_path}.zip" |
|
zip_directory(new_config.model_path, zip_path) |
|
entry = f"Fine-tuned Diffusion: {new_model_name}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Diffusion Model"), unsafe_allow_html=True) |
|
csv_path = f"sft_dataset_{int(time.time())}.csv" |
|
with open(csv_path, "w", newline="") as f: |
|
writer = csv.writer(f) |
|
writer.writerow(["image", "text"]) |
|
for _, row in edited_data.iterrows(): |
|
writer.writerow([row["image"], row["text"]]) |
|
st.markdown(get_download_link(csv_path, "text/csv", "Download SFT Dataset CSV"), unsafe_allow_html=True) |
|
|
|
with tab5: |
|
st.header("Test Titan 🧪") |
|
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False): |
|
st.warning("Please build or load a Titan first! ⚠️") |
|
else: |
|
if isinstance(st.session_state['builder'], ModelBuilder): |
|
if st.session_state['builder'].sft_data: |
|
st.write("Testing with SFT Data:") |
|
for item in st.session_state['builder'].sft_data[:3]: |
|
prompt = item["prompt"] |
|
expected = item["response"] |
|
status_container = st.empty() |
|
generated = st.session_state['builder'].evaluate(prompt, status_container) |
|
st.write(f"**Prompt**: {prompt}") |
|
st.write(f"**Expected**: {expected}") |
|
st.write(f"**Generated**: {generated}") |
|
st.write("---") |
|
status_container.empty() |
|
test_prompt = st.text_area("Enter Test Prompt", "What is AI?") |
|
if st.button("Run Test ▶️"): |
|
status_container = st.empty() |
|
result = st.session_state['builder'].evaluate(test_prompt, status_container) |
|
entry = f"Causal LM Test: {test_prompt} -> {result}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.write(f"**Generated Response**: {result}") |
|
status_container.empty() |
|
elif isinstance(st.session_state['builder'], DiffusionBuilder): |
|
test_prompt = st.text_area("Enter Test Prompt", "Neon Batman") |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
if st.button("Run Test ▶️"): |
|
image = st.session_state['builder'].generate(test_prompt) |
|
output_file = generate_filename("diffusion_test", "png") |
|
image.save(output_file) |
|
entry = f"Diffusion Test: {test_prompt} -> {output_file}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.image(image, caption="Generated Image") |
|
update_gallery() |
|
|
|
with tab6: |
|
st.header("Agentic RAG Party 🌐") |
|
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False): |
|
st.warning("Please build or load a Titan first! ⚠️") |
|
else: |
|
if isinstance(st.session_state['builder'], ModelBuilder): |
|
if st.button("Run NLP RAG Demo 🎉"): |
|
agent = PartyPlannerAgent(st.session_state['builder'].model, st.session_state['builder'].tokenizer) |
|
task = "Plan a luxury superhero-themed party at Wayne Manor." |
|
plan_df = agent.plan_party(task) |
|
entry = f"NLP RAG Demo: Planned party at Wayne Manor" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.dataframe(plan_df) |
|
elif isinstance(st.session_state['builder'], DiffusionBuilder): |
|
if st.button("Run CV RAG Demo 🎉"): |
|
agent = CVPartyPlannerAgent(st.session_state['builder'].pipeline) |
|
task = "Generate images for a luxury superhero-themed party." |
|
plan_df = agent.plan_party(task) |
|
entry = f"CV RAG Demo: Generated party images" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.dataframe(plan_df) |
|
for _, row in plan_df.iterrows(): |
|
image = agent.generate(row["Image Idea"]) |
|
output_file = generate_filename(f"cv_rag_{row['Theme'].lower()}", "png") |
|
image.save(output_file) |
|
st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}") |
|
update_gallery() |
|
|
|
with tab7: |
|
st.header("Test OCR 🔍") |
|
captured_files = get_gallery_files(["png"]) |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
all_files = captured_files + selected_pdfs |
|
if all_files: |
|
selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select") |
|
if selected_file: |
|
if selected_file.endswith('.png'): |
|
image = Image.open(selected_file) |
|
else: |
|
doc = fitz.open(selected_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() |
|
st.image(image, caption="Input Image", use_container_width=True) |
|
if st.button("Run OCR 🚀", key="ocr_run"): |
|
output_file = generate_filename("ocr_output", "txt") |
|
st.session_state['processing']['ocr'] = True |
|
result = asyncio.run(process_ocr(image, output_file)) |
|
entry = f"OCR Test: {selected_file} -> {output_file}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.text_area("OCR Result", result, height=200, key="ocr_result") |
|
st.success(f"OCR output saved to {output_file}") |
|
st.session_state['processing']['ocr'] = False |
|
else: |
|
st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!") |
|
|
|
with tab8: |
|
st.header("Test Image Gen 🎨") |
|
captured_files = get_gallery_files(["png"]) |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
all_files = captured_files + selected_pdfs |
|
if all_files: |
|
selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select") |
|
if selected_file: |
|
if selected_file.endswith('.png'): |
|
image = Image.open(selected_file) |
|
else: |
|
doc = fitz.open(selected_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() |
|
st.image(image, caption="Reference Image", use_container_width=True) |
|
prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt") |
|
if st.button("Run Image Gen 🚀", key="gen_run"): |
|
output_file = generate_filename("gen_output", "png") |
|
st.session_state['processing']['gen'] = True |
|
result = asyncio.run(process_image_gen(prompt, output_file)) |
|
entry = f"Image Gen Test: {prompt} -> {output_file}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.image(result, caption="Generated Image", use_container_width=True) |
|
st.success(f"Image saved to {output_file}") |
|
st.session_state['processing']['gen'] = False |
|
else: |
|
st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!") |
|
|
|
with tab9: |
|
st.header("Custom Diffusion 🎨🤓") |
|
st.write("Unleash your inner artist with our tiny diffusion models!") |
|
captured_files = get_gallery_files(["png"]) |
|
selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)] |
|
all_files = captured_files + selected_pdfs |
|
if all_files: |
|
st.subheader("Select Images or PDFs to Train") |
|
selected_files = st.multiselect("Pick Images or PDFs", all_files, key="diffusion_select") |
|
images = [] |
|
for file in selected_files: |
|
if file.endswith('.png'): |
|
images.append(Image.open(file)) |
|
else: |
|
doc = fitz.open(file) |
|
pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) |
|
images.append(Image.frombytes("RGB", [pix.width, pix.height], pix.samples)) |
|
doc.close() |
|
|
|
model_options = [ |
|
("PixelTickler 🎨✨", "OFA-Sys/small-stable-diffusion-v0"), |
|
("DreamWeaver 🌙🖌️", "stabilityai/stable-diffusion-2-base"), |
|
("TinyArtBot 🤖🖼️", "custom") |
|
] |
|
model_choice = st.selectbox("Choose Your Diffusion Dynamo", [opt[0] for opt in model_options], key="diffusion_model") |
|
model_name = next(opt[1] for opt in model_options if opt[0] == model_choice) |
|
|
|
if st.button("Train & Generate 🚀", key="diffusion_run"): |
|
output_file = generate_filename("custom_diffusion", "png") |
|
st.session_state['processing']['diffusion'] = True |
|
if model_name == "custom": |
|
result = asyncio.run(process_custom_diffusion(images, output_file, model_choice)) |
|
else: |
|
builder = DiffusionBuilder() |
|
builder.load_model(model_name) |
|
result = builder.generate("A superhero scene inspired by captured images") |
|
result.save(output_file) |
|
entry = f"Custom Diffusion: {model_choice} -> {output_file}" |
|
if entry not in st.session_state['history']: |
|
st.session_state['history'].append(entry) |
|
st.image(result, caption=f"{model_choice} Masterpiece", use_container_width=True) |
|
st.success(f"Image saved to {output_file}") |
|
st.session_state['processing']['diffusion'] = False |
|
else: |
|
st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!") |
|
|
|
update_gallery() |
