t

app.py

@@ -34,24 +34,20 @@ def process_pdf_gemma(pdf_path, model_name, progress=gr.Progress()):
     return gemma_handler.process_pdf(pdf_path, model_name, progress)
 
 @spaces.GPU
-def load_model_multiline(model_path):
-    return gemma_multiline_handler.load_model(model_path)
+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)
 
 @spaces.GPU
-def process_image_multiline(image, temp, top_p, repetition_penalty):
-    return gemma_multiline_handler.generate_text_from_image(image, temp, top_p, repetition_penalty)
+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)
 
 @spaces.GPU
-def process_image_multiline_stream(image, temp, top_p, repetition_penalty):
-    yield from gemma_multiline_handler.generate_text_stream(image, temp, top_p, repetition_penalty)
+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)
 
 @spaces.GPU
-def process_pdf_multiline(pdf, temp, top_p, repetition_penalty):
-    return gemma_multiline_handler.process_pdf(pdf, temp, top_p, repetition_penalty)
-
-@spaces.GPU
-def process_pdf_multiline_stream(pdf, temp, top_p, repetition_penalty):
-    yield from gemma_multiline_handler.process_pdf_stream(pdf, temp, top_p, repetition_penalty)
+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)
 
 # Example images for document-level OCR
 document_examples = [
@@ -144,12 +140,12 @@ with gr.Blocks(title="Dhivehi Image to Text",css=css) as demo:
 
                     generate_button.click(
                         fn=process_image_multiline,
-                        inputs=[image_input, temperature_slider, top_p_slider, repetition_penalty_slider],
+                        inputs=[model_path_dropdown, image_input, temperature_slider, top_p_slider, repetition_penalty_slider],
                         outputs=text_output
                     )
 
                     show_event = stream_button.click(fn=show_stop_button_image, outputs=[stop_button, stream_button, generate_button])
-                    gen_event = show_event.then(fn=process_image_multiline_stream, inputs=[image_input, temperature_slider, top_p_slider, repetition_penalty_slider], outputs=text_output)
+                    gen_event = show_event.then(fn=process_image_multiline_stream, inputs=[model_path_dropdown, image_input, temperature_slider, top_p_slider, repetition_penalty_slider], outputs=text_output)
                     gen_event.then(fn=hide_stop_button_image, outputs=[stop_button, stream_button, generate_button])
                     stop_button.click(fn=hide_stop_button_image, outputs=[stop_button, stream_button, generate_button], cancels=[gen_event])
 
@@ -186,16 +182,16 @@ with gr.Blocks(title="Dhivehi Image to Text",css=css) as demo:
 
                     pdf_generate_button.click(
                         fn=process_pdf_multiline,
-                        inputs=[pdf_input, temperature_slider, top_p_slider, repetition_penalty_slider],
+                        inputs=[model_path_dropdown, pdf_input, temperature_slider, top_p_slider, repetition_penalty_slider],
                         outputs=pdf_text_output
                     )
 
                     pdf_show_event = pdf_stream_button.click(fn=show_stop_button_pdf, outputs=[pdf_stop_button, pdf_stream_button, pdf_generate_button])
-                    pdf_gen_event = pdf_show_event.then(fn=process_pdf_multiline_stream, inputs=[pdf_input, temperature_slider, top_p_slider, repetition_penalty_slider], outputs=pdf_text_output)
+                    pdf_gen_event = pdf_show_event.then(fn=process_pdf_multiline_stream, inputs=[model_path_dropdown, pdf_input, temperature_slider, top_p_slider, repetition_penalty_slider], outputs=pdf_text_output)
                     pdf_gen_event.then(fn=hide_stop_button_pdf, outputs=[pdf_stop_button, pdf_stream_button, pdf_generate_button])
                     pdf_stop_button.click(fn=hide_stop_button_pdf, outputs=[pdf_stop_button, pdf_stream_button, pdf_generate_button], cancels=[pdf_gen_event])
 
-            model_path_dropdown.change(fn=load_model_multiline, inputs=model_path_dropdown)
+            # model_path_dropdown.change(fn=load_model_multiline, inputs=model_path_dropdown)
 
         with gr.Tab("PaliGemma"):
             model_dropdown_paligemma = gr.Dropdown(

gemma_multiline.py

@@ -17,14 +17,14 @@ class GemmaMultilineHandler:
     def __init__(self):
         self.model = None
         self.processor = None
-        self.current_model_path = None
+        self.current_model_name = None
         self.instruction = 'Extract the dhivehi text from the image'
 
     def load_model(self, model_name: str):
         if not model_name:
             self.model = None
             self.processor = None
-            self.current_model_path = None
+            self.current_model_name = None
             print("Model name is empty. No model loaded.")
             return
 
@@ -33,8 +33,8 @@ class GemmaMultilineHandler:
             print(f"Model '{model_name}' not found.")
             return
 
-        if model_path == self.current_model_path and self.model is not None:
-            print(f"Model from {model_path} is already loaded.")
+        if model_name == self.current_model_name and self.model is not None:
+            print(f"Model '{model_name}' is already loaded.")
             return
 
         try:
@@ -44,12 +44,12 @@ class GemmaMultilineHandler:
                 torch_dtype=torch.bfloat16,
             )
             self.processor = AutoProcessor.from_pretrained(model_path)
-            self.current_model_path = model_path
+            self.current_model_name = model_name
             print(f"Model loaded from {model_path}")
         except Exception as e:
             self.model = None
             self.processor = None
-            self.current_model_path = None
+            self.current_model_name = None
             print(f"Failed to load model: {e}")
 
     def process_vision_info(self, messages: list[dict]) -> list[Image.Image]:
@@ -65,9 +65,15 @@ class GemmaMultilineHandler:
                     image_inputs.append(image.convert("RGB"))
         return image_inputs
 
-    def generate_text_from_image(self, image: Image.Image, temperature: float = 0.8, top_p: float = 1.0, repetition_penalty: float = 1.2) -> str:
+    def generate_text_from_image(self, model_name: str, image: Image.Image, temperature: float = 0.8, top_p: float = 1.0, repetition_penalty: float = 1.2, progress=None) -> str:
+        if model_name != self.current_model_name:
+            try:
+                if progress: progress(0, desc=f"Loading {model_name}...")
+            except: pass
+            self.load_model(model_name)
+
         if self.model is None or self.processor is None:
-            return "Model not loaded. Please provide a model path."
+            return "Model not loaded. Please select a model."
 
         messages = [
             {
@@ -114,7 +120,13 @@ class GemmaMultilineHandler:
         )
         return output_text[0]
 
-    def generate_text_stream(self, image: Image.Image, temperature: float = 0.8, top_p: float = 1.0, repetition_penalty: float = 1.2):
+    def generate_text_stream(self, model_name: str, image: Image.Image, temperature: float = 0.8, top_p: float = 1.0, repetition_penalty: float = 1.2, progress=None):
+        if model_name != self.current_model_name:
+            try:
+                if progress: progress(0, desc=f"Loading {model_name}...")
+            except: pass
+            self.load_model(model_name)
+        
         if self.model is None or self.processor is None:
             yield "Model not loaded. Please provide a model path."
             return
@@ -173,7 +185,13 @@ class GemmaMultilineHandler:
                 clean_text = generated_text[response_start_index:]
                 yield clean_text.strip()
 
-    def process_pdf(self, pdf_path, temperature, top_p, repetition_penalty):
+    def process_pdf(self, model_name: str, pdf_path, temperature, top_p, repetition_penalty, progress=None):
+        if model_name != self.current_model_name:
+            try:
+                if progress: progress(0, desc=f"Loading {model_name}...")
+            except: pass
+            self.load_model(model_name)
+
         if self.model is None or self.processor is None:
             return "Model not loaded. Please load a model first."
         if pdf_path is None:
@@ -185,14 +203,20 @@ class GemmaMultilineHandler:
                 pix = page.get_pixmap()
                 image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                 doc.close()
-                return self.generate_text_from_image(image, temperature, top_p, repetition_penalty)
+                return self.generate_text_from_image(model_name, image, temperature, top_p, repetition_penalty, progress)
             else:
                 doc.close()
                 return "PDF has no pages."
         except Exception as e:
             return f"Failed to process PDF: {e}"
 
-    def process_pdf_stream(self, pdf_path, temperature, top_p, repetition_penalty):
+    def process_pdf_stream(self, model_name: str, pdf_path, temperature, top_p, repetition_penalty, progress=None):
+        if model_name != self.current_model_name:
+            try:
+                if progress: progress(0, desc=f"Loading {model_name}...")
+            except: pass
+            self.load_model(model_name)
+
         if self.model is None or self.processor is None:
             yield "Model not loaded. Please load a model first."
             return
@@ -206,7 +230,7 @@ class GemmaMultilineHandler:
                 pix = page.get_pixmap()
                 image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                 doc.close()
-                yield from self.generate_text_stream(image, temperature, top_p, repetition_penalty)
+                yield from self.generate_text_stream(model_name, image, temperature, top_p, repetition_penalty, progress)
             else:
                 doc.close()
                 yield "PDF has no pages."