Update app.py

app.py

@@ -2,12 +2,12 @@
 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
-import requests
 from PIL import Image
 import logging
 import asyncio
@@ -82,7 +82,7 @@ class DiffusionBuilder:
     def generate(self, prompt: str):
         return self.pipeline(prompt, num_inference_steps=20).images[0]
 
-def generate_filename(sequence, ext="png"):
+def generate_filename(sequence, ext):
     timestamp = time.strftime("%d%m%Y%H%M%S")
     return f"{sequence}_{timestamp}.{ext}"
 
@@ -98,35 +98,81 @@ async def process_image_gen(prompt, output_file, builder):
     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}"
+# Upload Functions
+def upload_images(files, history, selected_files):
+    if not files:
+        return "No files uploaded", history, selected_files
+    uploaded = []
+    for file in files:
+        ext = file.name.split('.')[-1].lower()
+        if ext in ["jpg", "png"]:
+            output_path = f"img_{int(time.time())}_{os.path.basename(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
+            uploaded.append(output_path)
+            history.append(f"Uploaded Image: {output_path}")
+            selected_files[output_path] = False
+    return f"Uploaded {len(uploaded)} images", history, selected_files
 
+def upload_videos(files, history, selected_files):
+    if not files:
+        return "No files uploaded", history, selected_files
+    uploaded = []
+    for file in files:
+        ext = file.name.split('.')[-1].lower()
+        if ext == "mp4":
+            output_path = f"vid_{int(time.time())}_{os.path.basename(file.name)}"
+            with open(output_path, "wb") as f:
+                f.write(file.read())
+            uploaded.append(output_path)
+            history.append(f"Uploaded Video: {output_path}")
+            selected_files[output_path] = False
+    return f"Uploaded {len(uploaded)} videos", history, selected_files
+
+def upload_documents(files, history, selected_files):
+    if not files:
+        return "No files uploaded", history, selected_files
+    uploaded = []
+    for file in files:
+        ext = file.name.split('.')[-1].lower()
+        if ext in ["md", "pdf", "docx"]:
+            output_path = f"doc_{int(time.time())}_{os.path.basename(file.name)}"
+            with open(output_path, "wb") as f:
+                f.write(file.read())
+            uploaded.append(output_path)
+            history.append(f"Uploaded Document: {output_path}")
+            selected_files[output_path] = False
+    return f"Uploaded {len(uploaded)} documents", history, selected_files
+
+def upload_datasets(files, history, selected_files):
+    if not files:
+        return "No files uploaded", history, selected_files
+    uploaded = []
+    for file in files:
+        ext = file.name.split('.')[-1].lower()
+        if ext in ["csv", "xlsx"]:
+            output_path = f"data_{int(time.time())}_{os.path.basename(file.name)}"
+            with open(output_path, "wb") as f:
+                f.write(file.read())
+            uploaded.append(output_path)
+            history.append(f"Uploaded Dataset: {output_path}")
+            selected_files[output_path] = False
+    return f"Uploaded {len(uploaded)} datasets", history, selected_files
+
+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
+    links = list(zip(links_title.split('\n'), links_url.split('\n')))
+    uploaded = []
+    for title, url in links:
+        if title and url:
+            link_entry = f"[{title}]({url})"
+            uploaded.append(link_entry)
+            history.append(f"Added Link: {link_entry}")
+            selected_files[link_entry] = False
+    return f"Added {len(uploaded)} links", history, selected_files
+
+# Gallery Update
 def update_galleries(history, selected_files):
     galleries = {
         "images": get_gallery_files(["jpg", "png"]),
@@ -136,15 +182,22 @@ def update_galleries(history, selected_files):
         "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]]
+        "images": [(Image.open(f), os.path.basename(f)) for f in galleries["images"]],
+        "videos": [(f, os.path.basename(f)) for f in galleries["videos"]],  # File path as preview
+        "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"]],
+        "datasets": [(f, os.path.basename(f)) for f in galleries["datasets"]],
+        "links": [(f, f.split(']')[0][1:]) for f in galleries["links"]]
     }
     history.append(f"Updated galleries: {sum(len(g) for g in galleries.values())} files")
     return gallery_outputs, history, selected_files
 
+# Sidebar Update
+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(')')]
+    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]
+    return file_list, history
+
+# Operations
 def toggle_selection(file_list, selected_files):
     for file in file_list:
         selected_files[file] = not selected_files.get(file, False)
@@ -169,76 +222,47 @@ with gr.Blocks(title="AI Vision & SFT Titans 🚀") as demo:
 
     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)
+            gr.Markdown("## 📁 Files")
+            sidebar_files = gr.Files(label="Downloads", height=300)
 
         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)
+                with gr.TabItem("📤 Upload"):
+                    with gr.Row():
+                        img_upload = gr.File(label="🖼️ Images (jpg/png)", file_count="multiple", accept=["image/jpeg", "image/png"])
+                        vid_upload = gr.File(label="🎥 Videos (mp4)", file_count="multiple", accept=["video/mp4"])
+                    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"])
+                        data_upload = gr.File(label="📊 Data (csv/xlsx)", file_count="multiple", accept=["text/csv", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"])
+                    with gr.Row():
+                        links_title = gr.Textbox(label="🔗 Link Titles", lines=3)
+                        links_url = gr.Textbox(label="🔗 Link URLs", lines=3)
                     upload_status = gr.Textbox(label="Status")
+                    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"):
+                    img_gallery = gr.Gallery(label="🖼️ Images (jpg/png)", columns=4, height="auto")
+                    vid_gallery = gr.Gallery(label="🎥 Videos (mp4)", columns=4, height="auto")
+                    doc_gallery = gr.Gallery(label="📜 Docs (md/pdf/docx)", columns=4, height="auto")
+                    data_gallery = gr.Gallery(label="📊 Data (csv/xlsx)", columns=4, height="auto")
+                    link_gallery = gr.Gallery(label="🔗 Links", columns=4, height="auto")
+                    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"):
                     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])
+                    select_files = gr.Dropdown(choices=list(selected_files.value.keys()), multiselect=True, label="Select Files")
+                    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])
 
 demo.launch()