Create app.py

app.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+import os
+import glob
+import time
+import pandas as pd
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from diffusers import StableDiffusionPipeline
+import fitz
+import requests
+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.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
+    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 get_gallery_files(file_types):
+    return sorted(list(set([f for ext in file_types for f in glob.glob(f"*.{ext}")])))  # Deduplicate files
+
+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
+
+# Smart Uploader Functions
+def upload_files(files, links_title, links_url, history, selected_files):
+    uploaded = {"images": [], "videos": [], "documents": [], "datasets": [], "links": []}
+    if files:
+        for file in files:
+            ext = file.name.split('.')[-1].lower()
+            output_path = f"uploaded_{int(time.time())}_{file.name}"
+            with open(output_path, "wb") as f:
+                f.write(file.read())
+            if ext in ["jpg", "png"]:
+                uploaded["images"].append(output_path)
+            elif ext == "mp4":
+                uploaded["videos"].append(output_path)
+            elif ext in ["md", "pdf", "docx"]:
+                uploaded["documents"].append(output_path)
+            elif ext in ["csv", "xlsx"]:
+                uploaded["datasets"].append(output_path)
+            history.append(f"Uploaded: {output_path}")
+            selected_files[output_path] = False  # Default unchecked
+    if links_title and links_url:
+        links = list(zip(links_title.split('\n'), links_url.split('\n')))
+        for title, url in links:
+            if title and url:
+                link_entry = f"[{title}]({url})"
+                uploaded["links"].append(link_entry)
+                history.append(f"Added Link: {link_entry}")
+                selected_files[link_entry] = False
+    return uploaded, history, selected_files
+
+def update_galleries(history, selected_files):
+    galleries = {
+        "images": get_gallery_files(["jpg", "png"]),
+        "videos": get_gallery_files(["mp4"]),
+        "documents": get_gallery_files(["md", "pdf", "docx"]),
+        "datasets": get_gallery_files(["csv", "xlsx"]),
+        "links": [f for f in selected_files.keys() if f.startswith('[') and '](' in f and f.endswith(')')]
+    }
+    gallery_outputs = {
+        "images": [(Image.open(f), os.path.basename(f)) for f in galleries["images"][:4]],
+        "videos": [(f, os.path.basename(f)) for f in galleries["videos"][:4]],  # Video preview as file path
+        "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"][:4]],
+        "datasets": [(f, os.path.basename(f)) for f in galleries["datasets"][:4]],  # Text preview
+        "links": [(f, f.split(']')[0][1:]) for f in galleries["links"][:4]]
+    }
+    history.append(f"Updated galleries: {sum(len(g) for g in galleries.values())} files")
+    return gallery_outputs, history, selected_files
+
+def toggle_selection(file_list, selected_files):
+    for file in file_list:
+        selected_files[file] = not selected_files.get(file, False)
+    return selected_files
+
+def image_gen(prompt, builder, history, selected_files):
+    selected = [f for f, sel in selected_files.items() if sel and f.endswith(('.jpg', '.png'))]
+    if not selected:
+        return "No images selected", None, history, selected_files
+    output_file = generate_filename("gen_output", "png")
+    gen_image = asyncio.run(process_image_gen(prompt, output_file, builder))
+    history.append(f"Image Gen: {prompt} -> {output_file}")
+    selected_files[output_file] = True
+    return f"Image saved to {output_file}", gen_image, history, selected_files
+
+# Gradio UI
+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)
+    selected_files = gr.State(value={})
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 📁 File Tree")
+            with gr.Accordion("🌳 Uploads", open=True):
+                with gr.Row():
+                    gr.Markdown("### 🖼️ Images (jpg/png)")
+                    img_gallery = gr.Gallery(label="Images", columns=4, height="auto")
+                with gr.Row():
+                    gr.Markdown("### 🎥 Videos (mp4)")
+                    vid_gallery = gr.Gallery(label="Videos", columns=4, height="auto")
+                with gr.Row():
+                    gr.Markdown("### 📜 Docs (md/pdf/docx)")
+                    doc_gallery = gr.Gallery(label="Documents", columns=4, height="auto")
+                with gr.Row():
+                    gr.Markdown("### 📊 Data (csv/xlsx)")
+                    data_gallery = gr.Gallery(label="Datasets", columns=4, height="auto")
+                with gr.Row():
+                    gr.Markdown("### 🔗 Links")
+                    link_gallery = gr.Gallery(label="Links", columns=4, height="auto")
+            gr.Markdown("## 📜 History")
+            history_output = gr.Textbox(label="Log", lines=5, interactive=False)
+
+        with gr.Column(scale=3):
+            with gr.Row():
+                gr.Markdown("## 🛠️ Toolbar")
+                upload_btn = gr.Button("📤 Upload")
+                select_btn = gr.Button("✅ Select")
+                gen_btn = gr.Button("🎨 Generate")
+
+            with gr.Tabs():
+                with gr.TabItem("📤 Smart Upload"):
+                    file_upload = gr.File(label="Upload Files", file_count="multiple", type="binary")
+                    links_title = gr.Textbox(label="Link Titles (one per line)", lines=3)
+                    links_url = gr.Textbox(label="Link URLs (one per line)", lines=3)
+                    upload_status = gr.Textbox(label="Status")
+
+                with gr.TabItem("🔍 Operations"):
+                    prompt = gr.Textbox(label="Image Gen Prompt", value="Generate a neon version")
+                    op_status = gr.Textbox(label="Status")
+                    op_output = gr.Image(label="Output")
+
+            upload_btn.click(
+                upload_files,
+                inputs=[file_upload, links_title, links_url, history, selected_files],
+                outputs=[upload_status, history, selected_files]
+            ).then(
+                update_galleries,
+                inputs=[history, selected_files],
+                outputs=[img_gallery, vid_gallery, doc_gallery, data_gallery, link_gallery, history, selected_files]
+            )
+
+            select_btn.click(
+                toggle_selection,
+                inputs=[gr.Dropdown(choices=list(selected_files.value.keys()), multiselect=True, label="Select Files"), selected_files],
+                outputs=[selected_files]
+            ).then(
+                update_galleries,
+                inputs=[history, selected_files],
+                outputs=[img_gallery, vid_gallery, doc_gallery, data_gallery, link_gallery, history, selected_files]
+            )
+
+            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]
+            )
+
+    # Update history output
+    demo.load(lambda h: "\n".join(h[-5:]), inputs=[history], outputs=[history_output])
+
+demo.launch()