Update app.py

app.py

@@ -1,288 +1,159 @@
-#!/usr/bin/env python
-# app.py
-
 import io
 import os
 import re
-import base64
-import glob
-import logging
-import random
-import shutil
-import time
-import zipfile
-import json
-import asyncio
-
-from pathlib import Path
 from datetime import datetime
-from typing import Any, List, Dict, Optional
+from collections import Counter
 
 import pandas as pd
-import pytz
 import streamlit as st
-import aiofiles
-import requests
-
-from PIL import Image, ImageDraw, UnidentifiedImageError
+from PIL import Image
 from reportlab.pdfgen import canvas
 from reportlab.lib.utils import ImageReader
-from reportlab.lib.pagesizes import letter
-import fitz  # PyMuPDF
-
-from huggingface_hub import InferenceClient
-from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError
-
-# Optional AI/ML imports
-try:
-    import torch
-    from transformers import (
-        AutoModelForCausalLM,
-        AutoTokenizer,
-        AutoProcessor,
-        AutoModelForVision2Seq,
-        pipeline
-    )
-    _transformers_available = True
-except ImportError:
-    _transformers_available = False
-
-try:
-    from diffusers import StableDiffusionPipeline
-    _diffusers_available = True
-except ImportError:
-    _diffusers_available = False
 
-# --- Page Configuration ---
+# --- App Configuration ----------------------------------
 st.set_page_config(
-    page_title="Vision & Layout Titans (HF) 🚀🖼️",
-    page_icon="🤖",
+    page_title="Image → PDF Comic Layout",
     layout="wide",
     initial_sidebar_state="expanded",
-    menu_items={
-        'Get Help': 'https://huggingface.co/docs',
-        'About': "Combined App: Image→PDF Layout + HF AI Tools 🌌"
-    }
 )
 
-# --- Logging Setup ---
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger(__name__)
-log_records: List[logging.LogRecord] = []
-class LogCaptureHandler(logging.Handler):
-    def emit(self, record):
-        log_records.append(record)
-logger.addHandler(LogCaptureHandler())
-
-# --- Constants & Defaults ---
-HF_TOKEN = os.getenv("HF_TOKEN")
-DEFAULT_PROVIDER = "hf-inference"
-FEATURED_MODELS_LIST = [
-    "meta-llama/Meta-Llama-3.1-8B-Instruct",
-    "mistralai/Mistral-7B-Instruct-v0.3",
-    "google/gemma-2-9b-it",
-    "Qwen/Qwen2-7B-Instruct",
-    "microsoft/Phi-3-mini-4k-instruct",
-    "HuggingFaceH4/zephyr-7b-beta",
-    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-    "HuggingFaceTB/SmolLM-1.7B-Instruct"
-]
-
-# --- Session State Initialization ---
-def _init_state(key: str, default: Any):
-    if key not in st.session_state:
-        st.session_state[key] = default
-
-for k, v in {
-    'layout_snapshots': [],
-    'layout_new_uploads': [],
-    'layout_last_capture': None,
-    'history': [],
-    'processing': {},
-    'asset_checkboxes': {},
-    'downloaded_pdfs': {},
-    'unique_counter': 0,
-    'cam0_file': None,
-    'cam1_file': None,
-    'characters': [],
-    'char_form_reset_key': 0,
-    'gallery_size': 10,
-    'hf_inference_client': None,
-    'hf_provider': DEFAULT_PROVIDER,
-    'hf_custom_key': "",
-    'hf_selected_api_model': FEATURED_MODELS_LIST[0],
-    'hf_custom_api_model': "",
-    'local_models': {},
-    'selected_local_model_path': None,
-    'gen_max_tokens': 512,
-    'gen_temperature': 0.7,
-    'gen_top_p': 0.95,
-    'gen_frequency_penalty': 0.0,
-    'gen_seed': -1
-}.items():
-    _init_state(k, v)
-
-# --- Utility Functions ---
-def generate_filename(seq: str, ext: str = "png") -> str:
-    ts = time.strftime('%Y%m%d_%H%M%S')
-    safe = re.sub(r'[^\w\-]+', '_', seq)
-    return f"{safe}_{ts}.{ext}"
-
-def clean_stem(fn: str) -> str:
-    return os.path.splitext(os.path.basename(fn))[0].replace('-', ' ').replace('_', ' ').title()
-
-def get_download_link(path: str, mime: str, label: str = "Download") -> str:
-    if not os.path.exists(path): return f"{label} (not found)"
-    data = open(path,'rb').read()
-    b64 = base64.b64encode(data).decode()
-    return f'<a href="data:{mime};base64,{b64}" download="{os.path.basename(path)}">{label}</a>'
-
-def get_gallery_files(types: List[str] = ['png','jpg','jpeg','pdf','md','txt']) -> List[str]:
-    files = set()
-    for ext in types:
-        files.update(glob.glob(f"*.{ext}"))
-        files.update(glob.glob(f"*.{ext.upper()}"))
-    return sorted(files)
-
-# Delete with rerun
-def delete_asset(path: str):
-    try:
-        os.remove(path)
-        st.session_state['asset_checkboxes'].pop(path, None)
-        if path in st.session_state['layout_snapshots']:
-            st.session_state['layout_snapshots'].remove(path)
-        st.toast(f"Deleted {os.path.basename(path)}", icon="✅")
-    except OSError as e:
-        st.error(f"Delete failed: {e}")
-    st.rerun()
+st.title("🖼️ Image → PDF • Full-Page & Custom Layout Generator")
+st.markdown(
+    "Upload images, filter by orientation, reorder visually, and generate a captioned PDF where each page matches its image's dimensions."
+)
 
-# Sidebar gallery updater
-def update_gallery():
-    st.sidebar.markdown("### Asset Gallery 📸📖")
-    files = get_gallery_files()
-    if not files:
-        st.sidebar.info("No assets.")
-        return
-    st.sidebar.caption(f"Found {len(files)} assets.")
-    for f in files[:st.session_state['gallery_size']]:
-        name = os.path.basename(f)
-        ext = os.path.splitext(f)[1].lower()
-        st.sidebar.markdown(f"**{name}**")
-        with st.sidebar.expander("Preview", expanded=False):
-            try:
-                if ext in ['.png','.jpg','.jpeg']:
-                    st.image(Image.open(f), use_container_width=True)
-                elif ext == '.pdf':
-                    doc = fitz.open(f)
-                    if doc.page_count:
-                        pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5,0.5))
-                        img = Image.frombytes('RGB',[pix.width,pix.height],pix.samples)
-                        st.image(img, use_container_width=True)
-                    doc.close()
-                else:
-                    txt = Path(f).read_text(errors='ignore')
-                    st.code(txt[:200]+'…')
-            except:
-                st.warning("Preview error")
-        c1,c2,c3 = st.sidebar.columns(3)
-        sel = st.session_state['asset_checkboxes'].get(f, False)
-        c1.checkbox("Select", value=sel, key=f"cb_{f}")
-        st.session_state['asset_checkboxes'][f] = st.session_state.get(f"cb_{f}")
-        mime = {'png':'image/png','jpg':'image/jpeg','jpeg':'image/jpeg','pdf':'application/pdf','md':'text/markdown','txt':'text/plain'}.get(ext[1:], 'application/octet-stream')
-        with open(f,'rb') as fp:
-            c2.download_button("📥", data=fp, file_name=name, mime=mime, key=f"dl_{f}")
-        c3.button("🗑️", key=f"del_{f}", on_click=delete_asset, args=(f,))
-        st.sidebar.markdown("---")
+# --- Sidebar: Page Settings -----------------------------
+st.sidebar.header("1️⃣ Page Aspect Ratio & Size")
+ratio_map = {
+    "4:3 (Landscape)": (4, 3),
+    "16:9 (Landscape)": (16, 9),
+    "1:1 (Square)": (1, 1),
+    "2:3 (Portrait)": (2, 3),
+    "9:16 (Portrait)": (9, 16),
+}
+ratio_choice = st.sidebar.selectbox(
+    "Preset Ratio", list(ratio_map.keys()) + ["Custom…"]
+)
+if ratio_choice != "Custom…":
+    rw, rh = ratio_map[ratio_choice]
+else:
+    rw = st.sidebar.number_input("Custom Width Ratio", min_value=1, value=4)
+    rh = st.sidebar.number_input("Custom Height Ratio", min_value=1, value=3)
+BASE_WIDTH_PT = st.sidebar.slider(
+    "Base Page Width (pt)", min_value=400, max_value=1200, value=800, step=100
+)
+page_width = BASE_WIDTH_PT
+page_height = int(BASE_WIDTH_PT * (rh / rw))
+st.sidebar.markdown(f"**Preview page size:** {page_width}×{page_height} pt")
+
+# --- Main: Upload, Filter & Reorder -------------------
+st.header("2️⃣ Upload, Filter & Reorder Images")
+uploaded = st.file_uploader(
+    "📂 Select PNG/JPG images", type=["png","jpg","jpeg"], accept_multiple_files=True
+)
+ordered_files = []
+if uploaded:
+    records = []
+    for idx, f in enumerate(uploaded):
+        im = Image.open(f)
+        w, h = im.size
+        ar = round(w / h, 2)
+        orient = "Square" if 0.9 <= ar <= 1.1 else ("Landscape" if ar > 1.1 else "Portrait")
+        records.append({
+            "filename": f.name,
+            "width": w,
+            "height": h,
+            "aspect_ratio": ar,
+            "orientation": orient,
+            "order": idx,
+        })
+    df = pd.DataFrame(records)
+    dims = st.sidebar.multiselect(
+        "Include orientations:", options=["Landscape","Portrait","Square"],
+        default=["Landscape","Portrait","Square"]
+    )
+    df = df[df["orientation"].isin(dims)].reset_index(drop=True)
 
-# --- PDF Snapshot & Generation ---
-async def process_pdf_snapshot(path: str, mode: str='single', resF: float=2.0) -> List[str]:
-    status = st.empty()
-    status.text("Snapshot start...")
-    out_files: List[str] = []
+    st.markdown("#### Image Metadata & Order")
+    st.dataframe(df.style.format({"aspect_ratio": "{:.2f}"}), use_container_width=True)
+    st.markdown("*Drag rows or edit the `order` column to set PDF page sequence.*")
     try:
-        doc = fitz.open(path)
-        mat = fitz.Matrix(resF,resF)
-        cnt = {'single':1,'twopage':2,'allpages':len(doc)}.get(mode,1)
-        for i in range(min(cnt,len(doc))):
-            s = time.time()
-            page = doc[i]
-            pix = page.get_pixmap(matrix=mat)
-            base = os.path.splitext(os.path.basename(path))[0]
-            fname = generate_filename(f"{base}_pg{i+1}_{mode}","png")
-            await asyncio.to_thread(pix.save, fname)
-            out_files.append(fname)
-            status.text(f"Saved {fname} ({int(time.time()-s)}s)")
-        doc.close()
-        status.success(f"Snapshot done: {len(out_files)} files")
-    except Exception as e:
-        status.error(f"Snapshot error: {e}")
-        for f in out_files:
-            if os.path.exists(f): os.remove(f)
-        out_files = []
-    return out_files
-
-from reportlab.lib.pagesizes import letter
-
-def make_image_sized_pdf(sources: List[Any]) -> Optional[bytes]:
-    # dedupe
-    seen, uniq = set(), []
-    for s in sources:
-        key = s if isinstance(s,str) else getattr(s,'name',None)
-        if key and key not in seen:
-            seen.add(key)
-            uniq.append(s)
-    if not uniq:
-        st.warning("No images for PDF")
-        return None
-    buf = io.BytesIO()
-    c = canvas.Canvas(buf, pagesize=letter)
-    status = st.empty()
-    for idx,s in enumerate(uniq,1):
-        try:
-            img = Image.open(s) if isinstance(s,str) else Image.open(s)
-            w,h = img.size
-            cap = 30
-            c.setPageSize((w,h+cap))
-            c.drawImage(ImageReader(img),0,cap,w,h,mask='auto')
-            cap_txt = clean_stem(s if isinstance(s,str) else s.name)
-            c.setFont('Helvetica',12)
-            c.drawCentredString(w/2,cap/2,cap_txt)
-            c.setFont('Helvetica',8)
-            c.drawRightString(w-10,10,str(idx))
-            c.showPage()
-            status.text(f"Page {idx}/{len(uniq)} added")
-        except Exception as e:
-            status.error(f"Error page {idx}: {e}")
+        edited = st.experimental_data_editor(df, num_rows="fixed", use_container_width=True)
+        ordered_df = edited
+    except Exception:
+        edited = st.data_editor(
+            df,
+            column_config={
+                "order": st.column_config.NumberColumn(
+                    "Order", min_value=0, max_value=len(df)-1
+                )
+            },
+            hide_index=True,
+            use_container_width=True,
+        )
+        ordered_df = edited.sort_values("order").reset_index(drop=True)
+
+    name2file = {f.name: f for f in uploaded}
+    ordered_files = [name2file[n] for n in ordered_df["filename"] if n in name2file]
+
+# --- Utility: Clean stems -------------------------------
+def clean_stem(fn: str) -> str:
+    stem = os.path.splitext(fn)[0]
+    return stem.replace("-", " ").replace("_", " ")
+
+# --- PDF Creation: Image-sized + Captions --------------
+def make_image_sized_pdf(images):
+    buffer = io.BytesIO()
+    c = canvas.Canvas(buffer)
+    for idx, f in enumerate(images, start=1):
+        im = Image.open(f)
+        iw, ih = im.size
+        # Caption height
+        cap_h = 20
+        page_w, page_h = iw, ih + cap_h
+        c.setPageSize((page_w, page_h))
+        # Draw image
+        c.drawImage(ImageReader(im), 0, cap_h, iw, ih, preserveAspectRatio=True, mask='auto')
+        # Draw caption
+        caption = clean_stem(f.name)
+        c.setFont("Helvetica", 12)
+        c.drawCentredString(page_w/2, cap_h/2, caption)
+        # Page number
+        c.setFont("Helvetica", 8)
+        c.drawRightString(page_w - 10, 10, str(idx))
+        c.showPage()
     c.save()
-    buf.seek(0)
-    return buf.getvalue()
-
-# --- HF Inference Client ---
-def get_hf_client() -> Optional[InferenceClient]:
-    provider = st.session_state['hf_provider']
-    token = st.session_state['hf_custom_key'].strip() or HF_TOKEN
-    if provider!='hf-inference' and not token:
-        st.error(f"Provider {provider} needs token")
-        return None
-    client = st.session_state['hf_inference_client']
-    if not client:
-        st.session_state['hf_inference_client'] = InferenceClient(token=token, provider=provider)
-    return st.session_state['hf_inference_client']
-
-# --- HF Processing ---
-def process_text_hf(text: str, prompt: str, use_api: bool) -> str:
-    stp = st.empty(); stp.text("Processing...")
-    msgs = [{"role":"system","content":"You are an assistant."},
-            {"role":"user","content":f"{prompt}\n\n{text}"}]
-    out = ""
-    if use_api:
-        client = get_hf_client()
-        if not client: return "Client error"
-        model = st.session_state['hf_custom_api_model'] or st.session_state['hf_selected_api_model']
+    buffer.seek(0)
+    return buffer.getvalue()
+
+# --- Generate & Download -------------------------------
+st.header("3️⃣ Generate & Download PDF with Captions")
+if st.button("🖋️ Generate Captioned PDF"):
+    if not ordered_files:
+        st.warning("Upload and reorder at least one image.")
+    else:
+        # Timestamp + weekday
+        now = datetime.now()
+        prefix = now.strftime("%Y-%m%d-%I%M%p") + "-" + now.strftime("%a").upper()
+        # Build filename from ordered stems
+        stems = [clean_stem(f.name) for f in ordered_files]
+        basename = " - ".join(stems)
+        fname = f"{prefix}-{basename}.pdf"
+        pdf_bytes = make_image_sized_pdf(ordered_files)
+        st.success(f"✅ PDF ready: **{fname}**")
+        st.download_button(
+            "⬇️ Download PDF", data=pdf_bytes,
+            file_name=fname, mime="application/pdf"
+        )
+        st.markdown("#### Preview of Page 1")
         try:
-            resp = client.chat_completion(
-                model=model,
-                messages=msgs,
-                max_tokens=st.session_state['gen_max_tokens'],
-                temperature=st.session
-]}]}
+            import fitz
+            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+            pix = doc.load_page(0).get_pixmap(matrix=fitz.Matrix(1.5,1.5))
+            st.image(pix.tobytes(), use_container_width=True)
+        except Exception:
+            st.info("Install `pymupdf` (`fitz`) for preview.")
+
+# --- Footer ------------------------------------------------
+st.sidebar.markdown("---")
+st.sidebar.markdown("Built by Aaron C. Wacker • Senior AI Engineer")