4

app.py

@@ -1,50 +1,81 @@
 import spaces
 import gradio as gr
-import numpy as np
-from paligemma2 import PaliGemma2Handler, MODELS as PALIGEMMA_MODELS
-from gemma import GemmaHandler, MODELS as GEMMA_MODELS
-from gemma_multiline import GemmaMultilineHandler, MODELS as GEMMA_MULTILINE_MODELS
 
-# Initialize model handlers
-paligemma_handler = PaliGemma2Handler()
-gemma_handler = GemmaHandler()
-gemma_multiline_handler = GemmaMultilineHandler()
+# Model names defined here to avoid importing torch/transformers at startup.
+# Heavy imports (torch, transformers, peft) are deferred until first GPU call.
+PALIGEMMA_MODELS = {
+    "Medium-14k, Single Line": {},
+    "Medium-16k, Single Line": {},
+    "Small, Single Line": {},
+}
+GEMMA_MODELS = {
+    "Gemma-3 10k": {},
+}
+GEMMA_MULTILINE_MODELS = {
+    "Gemma Multiline - no-format": "",
+    "Gemma Multiline - line": "",
+}
+
+_paligemma_handler = None
+_gemma_handler = None
+_gemma_multiline_handler = None
+
+def get_paligemma_handler():
+    global _paligemma_handler
+    if _paligemma_handler is None:
+        from paligemma2 import PaliGemma2Handler
+        _paligemma_handler = PaliGemma2Handler()
+    return _paligemma_handler
+
+def get_gemma_handler():
+    global _gemma_handler
+    if _gemma_handler is None:
+        from gemma import GemmaHandler
+        _gemma_handler = GemmaHandler()
+    return _gemma_handler
+
+def get_gemma_multiline_handler():
+    global _gemma_multiline_handler
+    if _gemma_multiline_handler is None:
+        from gemma_multiline import GemmaMultilineHandler
+        _gemma_multiline_handler = GemmaMultilineHandler()
+    return _gemma_multiline_handler
 
 @spaces.GPU
 def process_image_paligemma(model_name, image, progress=gr.Progress()):
     """Process a single image with PaliGemma2"""
-    return paligemma_handler.process_image(model_name, image, progress)
+    return get_paligemma_handler().process_image(model_name, image, progress)
 
 @spaces.GPU
 def process_image_gemma(model_name, image, progress=gr.Progress()):
     """Process a single image with Gemma"""
-    return gemma_handler.process_image(model_name, image, progress)
+    return get_gemma_handler().process_image(model_name, image, progress)
 
 @spaces.GPU
 def process_pdf_paligemma(pdf_path, model_name, progress=gr.Progress()):
     """Process a PDF file with PaliGemma2"""
-    return paligemma_handler.process_pdf(pdf_path, model_name, progress)
+    return get_paligemma_handler().process_pdf(pdf_path, model_name, progress)
 
 @spaces.GPU
 def process_pdf_gemma(pdf_path, model_name, progress=gr.Progress()):
     """Process a PDF file with Gemma"""
-    return gemma_handler.process_pdf(pdf_path, model_name, progress)
+    return get_gemma_handler().process_pdf(pdf_path, model_name, progress)
 
 @spaces.GPU
 def process_image_multiline(model_name, image, temp, top_p, repetition_penalty, progress=gr.Progress()):
-    return gemma_multiline_handler.generate_text_from_image(model_name, image, temp, top_p, repetition_penalty, progress)
+    return get_gemma_multiline_handler().generate_text_from_image(model_name, image, temp, top_p, repetition_penalty, progress)
 
 @spaces.GPU
 def process_image_multiline_stream(model_name, image, temp, top_p, repetition_penalty, progress=gr.Progress()):
-    yield from gemma_multiline_handler.generate_text_stream(model_name, image, temp, top_p, repetition_penalty, progress)
+    yield from get_gemma_multiline_handler().generate_text_stream(model_name, image, temp, top_p, repetition_penalty, progress)
 
 @spaces.GPU
 def process_pdf_multiline(model_name, pdf, temp, top_p, repetition_penalty, progress=gr.Progress()):
-    return gemma_multiline_handler.process_pdf(model_name, pdf, temp, top_p, repetition_penalty, progress)
+    return get_gemma_multiline_handler().process_pdf(model_name, pdf, temp, top_p, repetition_penalty, progress)
 
 @spaces.GPU
 def process_pdf_multiline_stream(model_name, pdf, temp, top_p, repetition_penalty, progress=gr.Progress()):
-    yield from gemma_multiline_handler.process_pdf_stream(model_name, pdf, temp, top_p, repetition_penalty, progress)
+    yield from get_gemma_multiline_handler().process_pdf_stream(model_name, pdf, temp, top_p, repetition_penalty, progress)
 
 # Example images for document-level OCR
 document_examples = [