fixed docker model issues

services/OCR_glm_service.py

@@ -1,14 +1,20 @@
 from transformers import AutoModelForImageTextToText, AutoProcessor
 import torch
+import os
 from pathlib import Path
 from helpers import get_project_root
 
 
 class OCR_Glm_Service:
     def __init__(self, ocr_path=None, device=None):
+        ROOT = get_project_root()
+        self.base_model_path = Path(os.getenv("MODEL_PATH", ROOT / "backend" / "models"))
+            
         if not ocr_path:
-            ROOT = get_project_root()
-            ocr_path = ROOT / "backend" / "models" / "GlmOcr"
+            ocr_path = self.base_model_path / "GlmOcr"
+        else:
+            ocr_path = Path(ocr_path)
+
         processor_path = ocr_path / "processor"
         model_path = ocr_path / "model"
 
@@ -49,13 +55,12 @@ class OCR_Glm_Service:
         return output_text
     
     def load_model(self):
-        ROOT = get_project_root()
-        GLMOCR_MODEL_DIR = ROOT / "backend" / "models" / "GlmOcr"
-        MODEL_PATH = "zai-org/GLM-OCR"
+        GLMOCR_MODEL_DIR = self.base_model_path / "GlmOcr"
+        DOWNLOAD_MODEL = "zai-org/GLM-OCR"
 
-        model = AutoModelForImageTextToText.from_pretrained(MODEL_PATH)
+        model = AutoModelForImageTextToText.from_pretrained(DOWNLOAD_MODEL)
         model.save_pretrained(GLMOCR_MODEL_DIR / "model")
-        processor = AutoProcessor.from_pretrained( MODEL_PATH)
+        processor = AutoProcessor.from_pretrained( DOWNLOAD_MODEL)
         processor.save_pretrained(GLMOCR_MODEL_DIR / "processor")
 
         print(f"Downloaded GLM OCR to: {GLMOCR_MODEL_DIR}")

services/OCR_japanese_service.py

@@ -2,14 +2,19 @@ from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoder
 from PIL import Image
 from pathlib import Path
 from helpers import get_project_root
+import os
 
 ROOT = get_project_root()
 
 class OCR_Japanese_Service:
     def __init__(self, ocr_path=None, device=None):
+        ROOT = get_project_root()
+        self.base_model_path = Path(os.getenv("MODEL_PATH", ROOT / "backend" / "models"))
+
         if not ocr_path:
-            ROOT = get_project_root()
-            ocr_path = ROOT / "backend" / "models" / "Kha-white"
+            ocr_path = self.base_model_path / "Kha-white"
+        else:
+            ocr_path = Path(ocr_path)
 
         processor_path = ocr_path / "processor"
         model_path = ocr_path / "model"
@@ -43,13 +48,12 @@ class OCR_Japanese_Service:
         return generated_text
 
     def load_model(self):
-        ROOT = get_project_root()
-        MODEL_PATH = "kha-white/manga-ocr-base"
-        JAPANESE_OCR_DIR = ROOT / "backend" / "models" / "Kha-white"
+        DOWNLOAD_MODEL = "kha-white/manga-ocr-base"
+        JAPANESE_OCR_DIR = self.base_model_path / "Kha-white"
 
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-        model = VisionEncoderDecoderModel.from_pretrained(MODEL_PATH)
-        processor = AutoImageProcessor.from_pretrained(MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(DOWNLOAD_MODEL)
+        model = VisionEncoderDecoderModel.from_pretrained(DOWNLOAD_MODEL)
+        processor = AutoImageProcessor.from_pretrained(DOWNLOAD_MODEL)
 
         tokenizer.save_pretrained(JAPANESE_OCR_DIR / "tokenizer")
         model.save_pretrained(JAPANESE_OCR_DIR / "model")

services/bubble_detector_kiuyha_service.py

@@ -3,24 +3,27 @@ from PIL import Image
 from helpers import get_project_root
 from huggingface_hub import hf_hub_download
 from pathlib import Path
+import os
 
 class Bubble_Detector_Kiuyha_Service:
     def __init__(self, path=None):
+        ROOT = get_project_root()
+        self.base_model_path = Path(os.getenv("MODEL_PATH", ROOT / "backend" / "models"))
+
         if not path:
-            ROOT = get_project_root()
-            path =  ROOT / "backend" / "models"
+            path = self.base_model_path / "GlmOcr"
+        else:
+            path = Path(path)
             
-        model_path = path / "kiuyha.pt"
-
-        if not model_path.exists():
-            print(f"Kiuyha model not found at {model_path}. Attempting to download")
+        if not self.base_model_path.exists():
+            print(f"Kiuyha model not found at {self.base_model_path}. Attempting to download")
             self.load_model()
 
-        if model_path.exists():
-            self.model = YOLO(model_path)
+        if self.base_model_path.exists():
+            self.model = YOLO(self.base_model_path)
             print("Loaded Bubble Detector Kiuyha")
         else:
-            raise FileNotFoundError(f"Error: Could not find or retrieve {model_path}")
+            raise FileNotFoundError(f"Error: Could not find or retrieve {self.base_model_path}")
           
     def predict(self, img_path, conf=0.2, iou=0.4, show_labels=True, show_conf=True, imgsz=640):
         results = self.model.predict(
@@ -62,9 +65,7 @@ class Bubble_Detector_Kiuyha_Service:
         return sorted_boxes
     
     def load_model(self):
-        ROOT = get_project_root()
-        model_dir = ROOT / "backend" / "models"
-        target_path = model_dir / "kiuyha.pt"
+        target_path = self.base_model_path / "kiuyha.pt"
 
         if target_path.exists():
             print(f"Kiuya Model already exists at {target_path}")
@@ -73,7 +74,7 @@ class Bubble_Detector_Kiuyha_Service:
         downloaded_path = hf_hub_download(
             repo_id="Kiuyha/Manga-Bubble-YOLO",
             filename="model.pt",
-            local_dir=model_dir
+            local_dir=self.base_model_path
         )
 
         final_path = Path(downloaded_path).rename(target_path)

services/translate_qwen_service.py

@@ -1,14 +1,16 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
 import torch
+import os
+from pathlib import Path
 from helpers import get_project_root
 
 class Translate_Qwen_Service:
     def __init__(self, path=None, device=None):
-        if not path:
-            ROOT = get_project_root()
-            path = ROOT / "backend" / "models" / "Qwen"
-
+        ROOT = get_project_root()
+        self.base_model_path = Path(os.getenv("MODEL_PATH", ROOT / "backend" / "models"))
+        
+        path = self.base_model_path / "Qwen"
         tokenizer_path = path / "tokenizer"
         model_path = path / "model"
 
@@ -74,12 +76,11 @@ class Translate_Qwen_Service:
             return {"error": "Invalid JSON", "raw": output_text}
     
     def load_model(self):
-        ROOT = get_project_root()
-        MODEL_PATH = "Qwen/Qwen2.5-7B-Instruct"
-        QWEN_DIR = ROOT / "backend" / "models" / "Qwen"
+        QWEN_DIR = self.base_model_path / "Qwen"
+        DOWNLOAD_MODEL = "Qwen/Qwen2.5-7B-Instruct"
 
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-        model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(DOWNLOAD_MODEL)
+        model = AutoModelForCausalLM.from_pretrained(DOWNLOAD_MODEL)
         tokenizer.save_pretrained(QWEN_DIR / "tokenizer")
         model.save_pretrained(QWEN_DIR / "model")