Update app.py

app.py

@@ -13,7 +13,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 from diffusers import StableDiffusionPipeline
 from torch.utils.data import Dataset, DataLoader
 import csv
-from pdf2image import convert_from_path
+import fitz  # PyMuPDF, pure Python library
 import requests
 from PIL import Image
 import cv2
@@ -342,6 +342,77 @@ def download_pdf(url, output_path):
         logger.error(f"Failed to download {url}: {e}")
     return False
 
+# Async Processing Functions
+async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
+    start_time = time.time()
+    status = st.empty()
+    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
+    try:
+        doc = fitz.open(pdf_path)
+        output_files = []
+        if mode == "thumbnail":
+            page = doc[0]
+            pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))  # 50% scale
+            output_file = generate_filename("thumbnail", "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(1.0, 1.0))  # Full scale
+                output_file = generate_filename(f"twopage_{i}", "png")
+                pix.save(output_file)
+                output_files.append(output_file)
+        doc.close()
+        elapsed = int(time.time() - start_time)
+        status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
+        update_gallery()
+        return output_files
+    except Exception as e:
+        status.error(f"Failed to process PDF: {str(e)}")
+        return []
+
+async def process_ocr(image, output_file):
+    start_time = time.time()
+    status = st.empty()
+    status.text("Processing GOT-OCR2_0... (0s)")
+    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()
+    result = model.chat(tokenizer, image, ocr_type='ocr')
+    elapsed = int(time.time() - start_time)
+    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
+    async with aiofiles.open(output_file, "w") as f:
+        await f.write(result)
+    update_gallery()
+    return result
+
+async def process_image_gen(prompt, output_file):
+    start_time = time.time()
+    status = st.empty()
+    status.text("Processing Image Gen... (0s)")
+    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]
+    elapsed = int(time.time() - start_time)
+    status.text(f"Image Gen completed in {elapsed}s!")
+    gen_image.save(output_file)
+    update_gallery()
+    return gen_image
+
+async def process_custom_diffusion(images, output_file, model_name):
+    start_time = time.time()
+    status = st.empty()
+    status.text(f"Training {model_name}... (0s)")
+    unet = TinyUNet()
+    diffusion = TinyDiffusion(unet)
+    diffusion.train(images)
+    gen_image = diffusion.generate()
+    upscaled_image = diffusion.upscale(gen_image, scale_factor=2)
+    elapsed = int(time.time() - start_time)
+    status.text(f"{model_name} completed in {elapsed}s!")
+    upscaled_image.save(output_file)
+    update_gallery()
+    return upscaled_image
+
 # Mock Search Tool for RAG
 def mock_search(query: str) -> str:
     if "superhero" in query.lower():
@@ -423,73 +494,6 @@ def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_
     flight_time = (actual_distance / cruising_speed_kmh) + 1.0
     return round(flight_time, 2)
 
-# Async Processing Functions
-async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
-    start_time = time.time()
-    status = st.empty()
-    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
-    try:
-        images = convert_from_path(pdf_path, dpi=200)
-        output_files = []
-        if mode == "thumbnail":
-            img = images[0].resize((int(images[0].width * 0.5), int(images[0].height * 0.5)), Image.Resampling.LANCZOS)
-            output_file = generate_filename("thumbnail", "png")
-            img.save(output_file)
-            output_files.append(output_file)
-        elif mode == "twopage":
-            for i in range(min(2, len(images))):
-                output_file = generate_filename(f"twopage_{i}", "png")
-                images[i].save(output_file)
-                output_files.append(output_file)
-        elapsed = int(time.time() - start_time)
-        status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
-        update_gallery()
-        return output_files
-    except Exception as e:
-        status.error(f"Failed to process PDF: {str(e)}. Install poppler-utils (e.g., 'sudo apt-get install poppler-utils' on Ubuntu) and ensure it's in PATH.")
-        return []
-
-async def process_ocr(image, output_file):
-    start_time = time.time()
-    status = st.empty()
-    status.text("Processing GOT-OCR2_0... (0s)")
-    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()
-    result = model.chat(tokenizer, image, ocr_type='ocr')
-    elapsed = int(time.time() - start_time)
-    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
-    async with aiofiles.open(output_file, "w") as f:
-        await f.write(result)
-    update_gallery()
-    return result
-
-async def process_image_gen(prompt, output_file):
-    start_time = time.time()
-    status = st.empty()
-    status.text("Processing Image Gen... (0s)")
-    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]
-    elapsed = int(time.time() - start_time)
-    status.text(f"Image Gen completed in {elapsed}s!")
-    gen_image.save(output_file)
-    update_gallery()
-    return gen_image
-
-async def process_custom_diffusion(images, output_file, model_name):
-    start_time = time.time()
-    status = st.empty()
-    status.text(f"Training {model_name}... (0s)")
-    unet = TinyUNet()
-    diffusion = TinyDiffusion(unet)
-    diffusion.train(images)
-    gen_image = diffusion.generate()
-    upscaled_image = diffusion.upscale(gen_image, scale_factor=2)
-    elapsed = int(time.time() - start_time)
-    status.text(f"{model_name} completed in {elapsed}s!")
-    upscaled_image.save(output_file)
-    update_gallery()
-    return upscaled_image
-
 # Main App
 st.title("AI Vision & SFT Titans 🚀")