Add global device definition for all models

app.py

@@ -48,6 +48,7 @@ class Config:
     SEGMENTATION_REGION_IOU = 0.3
     SEGMENTATION_LINE_OVERLAP_THRESHOLD = 0.5
     SEGMENTATION_REGION_OVERLAP_THRESHOLD = 0.5
+    DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     ALLOWED_SOURCES = ("https://astia.narc.fi, /tmp/gradio")
     
     # Model paths
@@ -121,7 +122,8 @@ class HTRPipeline:
                  segmentation_line_iou: float = 0.3,
                  segmentation_region_iou: float = 0.3,
                  segmentation_line_overlap_threshold: float = 0.5,
-                 segmentation_region_overlap_threshold: float = 0.5
+                 segmentation_region_overlap_threshold: float = 0.5,
+                 device: str = 'cuda:0'
                  ):
         """
         Initialize HTR pipeline with segmentation and recognition models.
@@ -147,9 +149,15 @@ class HTRPipeline:
                                               segmentation_line_iou,
                                               segmentation_region_iou,
                                               segmentation_line_overlap_threshold,
-                                              segmentation_region_overlap_threshold
+                                              segmentation_region_overlap_threshold,
+                                              device
                                               )
-        self.recognizer = self._init_recognizer(recognition_model_path, recognition_batch_size)
+        self.recognizer = self._init_recognizer(
+            recognition_model_path, 
+            recognition_batch_size,
+            device
+        )
+        
         self.plotter = PlotHTR()
         
         if self.segmenter is None or self.recognizer is None:
@@ -164,7 +172,8 @@ class HTRPipeline:
                         segmentation_line_iou: float,
                         segmentation_region_iou: float,
                         segmentation_line_overlap_threshold: float,
-                        segmentation_region_overlap_threshold: float
+                        segmentation_region_overlap_threshold: float,
+                        device: str
                         ) -> Optional[SegmentImage]:
         """
         Initialize document segmentation model.
@@ -179,6 +188,7 @@ class HTRPipeline:
             segmentation_region_iou: IoU threshold for merging overlapping region polygons
             segmentation_line_overlap_threshold: Area overlap ratio threshold for merging lines
             segmentation_region_overlap_threshold: Area overlap ratio threshold for merging regions
+            device: Device where the model is run (cpu, cuda:0)
         Returns:
             Initialized SegmentImage instance or None if initialization fails
         """
@@ -192,7 +202,8 @@ class HTRPipeline:
                 line_iou=segmentation_line_iou,
                 region_iou=segmentation_region_iou,
                 line_overlap_threshold=segmentation_line_overlap_threshold,
-                region_overlap_threshold=segmentation_region_overlap_threshold
+                region_overlap_threshold=segmentation_region_overlap_threshold,
+                device=device
             )
             logger.info("✓ Segmentation model initialized")
             return segmenter
@@ -200,7 +211,10 @@ class HTRPipeline:
             logger.error(f"Failed to initialize segmentation model: {e}")
             return None
     
-    def _init_recognizer(self, model_path: str, batch_size: int) -> Optional[TextRecognition]:
+    def _init_recognizer(self, 
+                         model_path: str, 
+                         batch_size: int,
+                         device: str) -> Optional[TextRecognition]:
         """
         Initialize text recognition model.
         
@@ -214,7 +228,7 @@ class HTRPipeline:
         try:
             recognizer = TextRecognition(
                 model_path=model_path,
-                device='cuda:0' if torch.cuda.is_available() else 'cpu',
+                device=device,
                 batch_size=batch_size
             )
             logger.info("✓ Text recognition model initialized")
@@ -369,7 +383,8 @@ try:
         segmentation_line_iou = Config.SEGMENTATION_LINE_IOU,
         segmentation_region_iou = Config.SEGMENTATION_REGION_IOU,
         segmentation_line_overlap_threshold = Config.SEGMENTATION_LINE_OVERLAP_THRESHOLD,
-        segmentation_region_overlap_threshold = Config.SEGMENTATION_REGION_OVERLAP_THRESHOLD
+        segmentation_region_overlap_threshold = Config.SEGMENTATION_REGION_OVERLAP_THRESHOLD,
+        device = Config.DEVICE 
     )
     logger.info("✓ HTR Pipeline initialized successfully")
 except Exception as e: