app.py

#!/usr/bin/env python3
# 😂 Shebangin’ it like it’s 1999—Python 3, let’s roll!

# 🧳 Importing the whole circus—get ready for a wild ride!
import os
import glob
import time
import base64
import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from diffusers import StableDiffusionPipeline
import fitz
from PIL import Image
import logging
import asyncio
import aiofiles
from io import BytesIO
from dataclasses import dataclass
from typing import Optional
import gradio as gr

# 📜 Logging setup—because even AIs need a diary!
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
log_records = []  # 🗒️ Dear diary, today I logged a thing...

# 🤓 LogCaptureHandler class—catching logs like a pro fisherman!
class LogCaptureHandler(logging.Handler):
    # 🎣 Hooking those logs right outta the stream!
    def emit(self, record):
        log_records.append(record)

logger.addHandler(LogCaptureHandler())  # 🐟 Adding the hook to the logger—catch ‘em all!

# 🏰 ModelConfig dataclass—building castles for our AI kings!
@dataclass
class ModelConfig:
    name: str  # 🏷️ Naming our royal model
    base_model: str  # 🏛️ The foundation it stands on
    size: str  # 📏 How big’s this beast?
    domain: Optional[str] = None  # 🌍 Where does it rule? Optional kingdom!
    model_type: str = "causal_lm"  # ⚙️ What kind of magic does it wield?
    # 🗺️ Property to map the path—where the king resides!
    @property
    def model_path(self):
        return f"models/{self.name}"

# 🎨 DiffusionConfig dataclass—art school for diffusion models!
@dataclass
class DiffusionConfig:
    name: str  # 🖌️ What’s this masterpiece called?
    base_model: str  # 🖼️ The canvas it starts with
    size: str  # 📐 Size of the artwork, big or small?
    domain: Optional[str] = None  # 🎨 Optional style—abstract or realism?
    # 🗺️ Property to find the gallery—where the art hangs!
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

# 🤖 ModelBuilder class—assembling AI like Lego bricks!
class ModelBuilder:
    # 🛠️ Init—setting up the workshop!
    def __init__(self):
        self.config = None  # 🗺️ Blueprint? Not yet!
        self.model = None  # 🤖 The robot’s still in pieces
        self.tokenizer = None  # 📝 No word-chopper yet
    # 🚀 Load_model—blast off with a pre-trained brain!
    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
        self.model = AutoModelForCausalLM.from_pretrained(model_path)  # 🧠 Brain downloaded!
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)  # ✂️ Word-slicer ready!
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token  # 🩹 Patching up the tokenizer
        if config:
            self.config = config  # 📜 Got the blueprint now!
        self.model.to("cuda" if torch.cuda.is_available() else "cpu")  # ⚡ GPU or bust!
        return self
    # 💾 Save_model—stashing the AI for later glory!
    def save_model(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)  # 🏠 Making room for the save
        self.model.save_pretrained(path)  # 🧠 Brain archived!
        self.tokenizer.save_pretrained(path)  # ✂️ Slicer stored!

# 🎨 DiffusionBuilder class—crafting diffusion dreams one pixel at a time!
class DiffusionBuilder:
    # 🖌️ Init—prepping the easel for some art!
    def __init__(self):
        self.config = None  # 🗺️ No art plan yet
        self.pipeline = None  # 🎨 No paintbrush in hand
    # 🖼️ Load_model—grabbing a pre-painted masterpiece!
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")  # 🖌️ Brush loaded, CPU style!
        if config:
            self.config = config  # 📜 Art plan acquired!
        return self
    # 💾 Save_model—framing the artwork for the gallery!
    def save_model(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)  # 🏠 Prepping the gallery wall
        self.pipeline.save_pretrained(path)  # 🖼️ Hung up for all to see!
    # 🌈 Generate—spinning pixels into gold, Picasso-style!
    def generate(self, prompt: str):
        return self.pipeline(prompt, num_inference_steps=20).images[0]  # 🎨 Art in 20 steps or less!

# 😂 Time to stamp files like a boss—unique names incoming!
def generate_filename(sequence, ext):
    timestamp = time.strftime("%d%m%Y%H%M%S")  # ⏰ Clock says “name me now!”
    return f"{sequence}_{timestamp}.{ext}"

# 🕵️‍♂️ Sherlocking the filesystem for your precious files!
def get_gallery_files(file_types):
    return sorted(list(set([f for ext in file_types for f in glob.glob(f"*.{ext}")])))  # 🗃️ Deduped treasure hunt!

# 🎨 Paint the town neon—async image gen magic ahead!
async def process_image_gen(prompt, output_file, builder):
    if builder and isinstance(builder, DiffusionBuilder) and builder.pipeline:
        pipeline = builder.pipeline  # 🖌️ Using the pro’s brush!
    else:
        pipeline = StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu")  # 🎨 Default brush, CPU vibes!
    gen_image = pipeline(prompt, num_inference_steps=20).images[0]  # 🌟 20 steps to brilliance!
    gen_image.save(output_file)  # 🖼️ Saving the neon dream!
    return gen_image

# 🖼️ Snap those pics like a paparazzi—upload images with flair!
def upload_images(files, history, selected_files):
    if not files:
        return "No files uploaded", history, selected_files  # 😢 No pics, no party!
    uploaded = []
    for file in files:
        ext = file.name.split('.')[-1].lower()  # 🕵️ Sniffing out the file type!
        if ext in ["jpg", "png"]:
            output_path = f"img_{int(time.time())}_{os.path.basename(file.name)}"  # 🏷️ Tagging it fresh!
            with open(output_path, "wb") as f:
                f.write(file.read())  # 📸 Snap saved!
            uploaded.append(output_path)
            history.append(f"Uploaded Image: {output_path}")  # 📜 Logging the fame!
            selected_files[output_path] = False  # ✅ Unchecked for now!
    return f"Uploaded {len(uploaded)} images", history, selected_files

# 🎥 Roll camera—video uploads that’ll make Spielberg jealous!
def upload_videos(files, history, selected_files):
    if not files:
        return "No files uploaded", history, selected_files  # 🎬 No footage, no Oscar!
    uploaded = []
    for file in files:
        ext = file.name.split('.')[-1].lower()  # 🕵️ Checking the reel type!
        if ext == "mp4":
            output_path = f"vid_{int(time.time())}_{os.path.basename(file.name)}"  # 🎞️ Reel name ready!
            with open(output_path, "wb") as f:
                f.write(file.read())  # 🎥 Action, saved!
            uploaded.append(output_path)
            history.append(f"Uploaded Video: {output_path}")  # 📜 Cue the credits!
            selected_files[output_path] = False  # ✅ Not selected yet!
    return f"Uploaded {len(uploaded)} videos", history, selected_files

# 📜 Scribble some docs—PDFs and more, oh what a bore!
def upload_documents(files, history, selected_files):
    if not files:
        return "No files uploaded", history, selected_files  # 📝 No docs, no drama!
    uploaded = []
    for file in files:
        ext = file.name.split('.')[-1].lower()  # 🕵️ Peeking at the paper type!
        if ext in ["md", "pdf", "docx"]:
            output_path = f"doc_{int(time.time())}_{os.path.basename(file.name)}"  # 🏷️ Stamping the scroll!
            with open(output_path, "wb") as f:
                f.write(file.read())  # 📜 Scroll secured!
            uploaded.append(output_path)
            history.append(f"Uploaded Document: {output_path}")  # 📜 Noted in history!
            selected_files[output_path] = False  # ✅ Still on the bench!
    return f"Uploaded {len(uploaded)} documents", history, selected_files

# 📊 Data nerd alert—CSV and Excel uploads for the win!
def upload_datasets(files, history, selected_files):
    if not files:
        return "No files uploaded", history, selected_files  # 📈 No data, no geek-out!
    uploaded = []
    for file in files:
        ext = file.name.split('.')[-1].lower()  # 🕵️ Cracking the data code!
        if ext in ["csv", "xlsx"]:
            output_path = f"data_{int(time.time())}_{os.path.basename(file.name)}"  # 🏷️ Labeling the stats!
            with open(output_path, "wb") as f:
                f.write(file.read())  # 📊 Stats stashed!
            uploaded.append(output_path)
            history.append(f"Uploaded Dataset: {output_path}")  # 📜 Data’s in the books!
            selected_files[output_path] = False  # ✅ Not picked yet!
    return f"Uploaded {len(uploaded)} datasets", history, selected_files

# 🔗 Link it up—URLs and titles, the web’s wild child!
def upload_links(links_title, links_url, history, selected_files):
    if not links_title or not links_url:
        return "No links provided", history, selected_files  # 🌐 No links, no surf!
    links = list(zip(links_title.split('\n'), links_url.split('\n')))  # 🧩 Pairing titles and URLs!
    uploaded = []
    for title, url in links:
        if title and url:
            link_entry = f"[{title}]({url})"  # 🔗 Crafting the web gem!
            uploaded.append(link_entry)
            history.append(f"Added Link: {link_entry}")  # 📜 Linking history!
            selected_files[link_entry] = False  # ✅ Surf’s not up yet!
    return f"Added {len(uploaded)} links", history, selected_files

# 🖼️ Gallery glow-up—show off all your files in style!
def update_galleries(history, selected_files):
    galleries = {
        "images": get_gallery_files(["jpg", "png"]),  # 🖼️ Picture parade!
        "videos": get_gallery_files(["mp4"]),  # 🎥 Video vault!
        "documents": get_gallery_files(["md", "pdf", "docx"]),  # 📜 Doc depot!
        "datasets": get_gallery_files(["csv", "xlsx"]),  # 📊 Data den!
        "links": [f for f in selected_files.keys() if f.startswith('[') and '](' in f and f.endswith(')')]  # 🔗 Link lounge!
    }
    gallery_outputs = {
        "images": [(Image.open(f), os.path.basename(f)) for f in galleries["images"]],  # 🖼️ Picture perfect!
        "videos": [(f, os.path.basename(f)) for f in galleries["videos"]],  # 🎥 Reel deal!
        "documents": [(Image.frombytes("RGB", fitz.open(f)[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)).size, fitz.open(f)[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)).samples) if f.endswith('.pdf') else f, os.path.basename(f)) for f in galleries["documents"]],  # 📜 Doc dazzle!
        "datasets": [(f, os.path.basename(f)) for f in galleries["datasets"]],  # 📊 Data delight!
        "links": [(f, f.split(']')[0][1:]) for f in galleries["links"]]  # 🔗 Link love!
    }
    history.append(f"Updated galleries: {sum(len(g) for g in galleries.values())} files")  # 📜 Gallery grand total!
    return gallery_outputs, history, selected_files

# 📂 Sidebar swagger—download links that scream “take me home!”
def update_sidebar(history, selected_files):
    all_files = get_gallery_files(["jpg", "png", "mp4", "md", "pdf", "docx", "csv", "xlsx"]) + [f for f in selected_files.keys() if f.startswith('[') and '](' in f and f.endswith(')')]  # 🗃️ All the loot!
    file_list = [gr.File(label=os.path.basename(f) if not f.startswith('[') else f.split(']')[0][1:], value=f) for f in all_files]  # 📥 Download goodies!
    return file_list, history

# ✅ Check it or wreck it—toggle those selections like a pro!
def toggle_selection(file_list, selected_files):
    for file in file_list:
        selected_files[file] = not selected_files.get(file, False)  # ✅ Flip the switch, baby!
    return selected_files

# 🎨 Neon dreams unleashed—generate images that pop!
def image_gen(prompt, builder, history, selected_files):
    selected = [f for f, sel in selected_files.items() if sel and f.endswith(('.jpg', '.png'))]  # 🖼️ Picking the canvas!
    if not selected:
        return "No images selected", None, history, selected_files  # 😢 No art, no party!
    output_file = generate_filename("gen_output", "png")  # 🏷️ New masterpiece name!
    gen_image = asyncio.run(process_image_gen(prompt, output_file, builder))  # 🌈 Paint it neon!
    history.append(f"Image Gen: {prompt} -> {output_file}")  # 📜 Art history in the making!
    selected_files[output_file] = True  # ✅ Auto-select the new star!
    return f"Image saved to {output_file}", gen_image, history, selected_files

# 🎪 Gradio UI—step right up to the AI circus!
with gr.Blocks(title="AI Vision & SFT Titans 🚀") as demo:
    gr.Markdown("# AI Vision & SFT Titans 🚀")  # 🎉 Welcome to the big top!
    history = gr.State(value=[])  # 📜 The ringmaster’s logbook!
    builder = gr.State(value=None)  # 🤖 The AI acrobat, waiting in the wings!
    selected_files = gr.State(value={})  # ✅ The chosen ones, ready to perform!

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 📁 Files")  # 🗃️ The file circus tent!
            sidebar_files = gr.Files(label="Downloads", height=300)  # 📥 Grab your souvenirs here!

        with gr.Column(scale=3):
            with gr.Row():
                gr.Markdown("## 🛠️ Toolbar")  # 🔧 The circus control panel!
                select_btn = gr.Button("✅ Select")  # ✅ Pick your performers!
                gen_btn = gr.Button("🎨 Generate")  # 🎨 Unleash the art clowns!

            with gr.Tabs():
                with gr.TabItem("📤 Upload"):  # 📤 The upload trapeze!
                    with gr.Row():
                        img_upload = gr.File(label="🖼️ Images (jpg/png)", file_count="multiple", accept=["image/jpeg", "image/png"])  # 🖼️ Picture trapeze!
                        vid_upload = gr.File(label="🎥 Videos (mp4)", file_count="multiple", accept=["video/mp4"])  # 🎥 Video vault!
                    with gr.Row():
                        doc_upload = gr.File(label="📜 Docs (md/pdf/docx)", file_count="multiple", accept=["text/markdown", "application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"])  # 📜 Doc drop!
                        data_upload = gr.File(label="📊 Data (csv/xlsx)", file_count="multiple", accept=["text/csv", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"])  # 📊 Data dive!
                    with gr.Row():
                        links_title = gr.Textbox(label="🔗 Link Titles", lines=3)  # 🔗 Title tightrope!
                        links_url = gr.Textbox(label="🔗 Link URLs", lines=3)  # 🔗 URL unicycle!
                    upload_status = gr.Textbox(label="Status")  # 📢 Ringmaster’s update!
                    gr.Button("📤 Upload Images").click(upload_images, inputs=[img_upload, history, selected_files], outputs=[upload_status, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])
                    gr.Button("📤 Upload Videos").click(upload_videos, inputs=[vid_upload, history, selected_files], outputs=[upload_status, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])
                    gr.Button("📤 Upload Docs").click(upload_documents, inputs=[doc_upload, history, selected_files], outputs=[upload_status, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])
                    gr.Button("📤 Upload Data").click(upload_datasets, inputs=[data_upload, history, selected_files], outputs=[upload_status, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])
                    gr.Button("📤 Upload Links").click(upload_links, inputs=[links_title, links_url, history, selected_files], outputs=[upload_status, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), gr.Gallery(), history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])

                with gr.TabItem("🖼️ Gallery"):  # 🖼️ The big top showcase!
                    img_gallery = gr.Gallery(label="🖼️ Images (jpg/png)", columns=4, height="auto")  # 🖼️ Picture parade!
                    vid_gallery = gr.Gallery(label="🎥 Videos (mp4)", columns=4, height="auto")  # 🎥 Video vault!
                    doc_gallery = gr.Gallery(label="📜 Docs (md/pdf/docx)", columns=4, height="auto")  # 📜 Doc depot!
                    data_gallery = gr.Gallery(label="📊 Data (csv/xlsx)", columns=4, height="auto")  # 📊 Data den!
                    link_gallery = gr.Gallery(label="🔗 Links", columns=4, height="auto")  # 🔗 Link lounge!
                    gr.Button("🔄 Refresh").click(update_galleries, inputs=[history, selected_files], outputs=[img_gallery, vid_gallery, doc_gallery, data_gallery, link_gallery, history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])

                with gr.TabItem("🔍 Operations"):  # 🔍 The magic trick tent!
                    prompt = gr.Textbox(label="Image Gen Prompt", value="Generate a neon version")  # 🎨 Art spellbook!
                    op_status = gr.Textbox(label="Status")  # 📢 Trick status!
                    op_output = gr.Image(label="Output")  # 🎨 The big reveal!
                    select_files = gr.Dropdown(choices=list(selected_files.value.keys()), multiselect=True, label="Select Files")  # ✅ Pick your props!
                    select_btn.click(toggle_selection, inputs=[select_files, selected_files], outputs=[selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])
                    gen_btn.click(image_gen, inputs=[prompt, builder, history, selected_files], outputs=[op_status, op_output, history, selected_files]).then(update_galleries, inputs=[history, selected_files], outputs=[img_gallery, vid_gallery, doc_gallery, data_gallery, link_gallery, history, selected_files]).then(update_sidebar, inputs=[history, selected_files], outputs=[sidebar_files, history])

# 🎉 Launch the circus—step right up, folks!
demo.launch()