Update app.py

app.py

@@ -3,7 +3,6 @@ import os
 import glob
 import base64
 import time
-import shutil
 import pandas as pd
 import torch
 import torch.nn as nn
@@ -15,7 +14,6 @@ import csv
 import fitz
 import requests
 from PIL import Image
-import cv2
 import numpy as np
 import logging
 import asyncio
@@ -39,7 +37,6 @@ class LogCaptureHandler(logging.Handler):
 
 logger.addHandler(LogCaptureHandler())
 
-# Data Classes and Models (unchanged from your original code)
 @dataclass
 class ModelConfig:
     name: str
@@ -61,106 +58,12 @@ class DiffusionConfig:
     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 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())
-
-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 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):
         self.model = AutoModelForCausalLM.from_pretrained(model_path)
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
@@ -170,42 +73,10 @@ class ModelBuilder:
             self.config = config
         self.model.to("cuda" if torch.cuda.is_available() else "cpu")
         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):
-            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()
-            logger.info(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
-        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)
-    def evaluate(self, prompt: str):
-        self.model.eval()
-        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)
 
 class DiffusionBuilder:
     def __init__(self):
@@ -216,12 +87,14 @@ class DiffusionBuilder:
         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]
 
-# Utility Functions
 def generate_filename(sequence, ext="png"):
-    timestamp = time.strftime("%d%m%Y%HM%S")
+    timestamp = time.strftime("%d%m%Y%H%M%S")
     return f"{sequence}_{timestamp}.{ext}"
 
 def pdf_url_to_filename(url):
@@ -231,6 +104,11 @@ def pdf_url_to_filename(url):
 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}")])))  # Deduplicate files
 
+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)
@@ -252,25 +130,72 @@ async def process_pdf_snapshot(pdf_path, mode="single"):
         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
+
 # Gradio Interface Functions
-def update_gallery(history):
+def update_gallery(history, asset_checkboxes):
     all_files = get_gallery_files()
-    gallery_content = "\n".join([f"- {f}" for f in all_files[:5]])
+    gallery_images = []
+    for file in all_files[:5]:  # Limit to 5 for display
+        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_content, history
+    return gallery_images, history, asset_checkboxes
 
-def camera_snap(image, history):
+def camera_snap(image, cam_id, history, asset_checkboxes, cam_files):
     if image is not None:
-        filename = generate_filename("cam")
+        filename = generate_filename(f"cam{cam_id}")
         image.save(filename)
-        history.append(f"Snapshot saved: {filename}")
-        return f"Image saved as {filename}", history
-    return "No image captured", history
+        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):
+def download_pdfs(urls, history, asset_checkboxes):
     urls = urls.strip().split("\n")
     downloaded = []
     for url in urls:
@@ -279,7 +204,71 @@ def download_pdfs(urls, history):
             if download_pdf(url, output_path):
                 downloaded.append(output_path)
                 history.append(f"Downloaded PDF: {output_path}")
-    return f"Downloaded {len(downloaded)} PDFs", history
+                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)
@@ -289,62 +278,94 @@ def build_model(model_type, base_model, model_name, domain, history):
     history.append(f"Built {model_type} model: {model_name}")
     return builder, f"Model saved to {config.model_path}", history
 
-def test_model(builder, prompt, history):
-    if builder is None:
-        return "No model loaded", history
-    if isinstance(builder, ModelBuilder):
-        result = builder.evaluate(prompt)
-        history.append(f"Tested Causal LM: {prompt} -> {result}")
-        return result, history
-    elif isinstance(builder, DiffusionBuilder):
-        image = builder.generate(prompt)
-        output_file = generate_filename("diffusion_test")
-        image.save(output_file)
-        history.append(f"Tested Diffusion: {prompt} -> {output_file}")
-        return output_file, 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
 
 # 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)
+    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.Textbox(label="Gallery", lines=5)
-            gr.Button("Update Gallery").click(update_gallery, inputs=[history], outputs=[gallery_output, history])
-        
+            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 📷"):
-                    camera_input = gr.Image(type="pil", label="Take a Picture")
-                    snap_output = gr.Textbox(label="Status")
-                    gr.Button("Capture").click(camera_snap, inputs=[camera_input, history], outputs=[snap_output, history])
+                    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")
-                    gr.Button("Download").click(download_pdfs, inputs=[url_input, history], outputs=[pdf_output, history])
+                    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"] if model_type.value == "Causal LM" else ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"],
-                        label="Base Model"
+                        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="Domain", value="general")
+                    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 Titan 🧪"):
-                    test_prompt = gr.Textbox(label="Test Prompt", value="What is AI?")
-                    test_output = gr.Textbox(label="Result")
-                    gr.Button("Test").click(test_model, inputs=[builder, test_prompt, history], outputs=[test_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])
 
-    with gr.Row():
-        gr.Markdown("## History 📜")
-        history_output = gr.Textbox(value="\n".join(history.value), label="History", lines=5, interactive=False)
+    # Update history output on every interaction
+    demo.load(lambda h: "\n".join(h[-5:]), inputs=[history], outputs=[history_output])
 
 demo.launch()