Update app.py

app.py

@@ -6,7 +6,6 @@ from torchvision import transforms
 import numpy as np 
 from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
 import os
 import torch
 from pypinyin import pinyin
@@ -15,13 +14,68 @@ import json
 import os
 import torchvision.transforms.functional as F
 from statistics import mean
-
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize LLM as None - will be loaded lazily
+llm = None
+
+def load_llm():
+    """Lazy loading of LLM to avoid startup delays"""
+    global llm
+    if llm is None:
+        try:
+            logger.info("Loading LLM model...")
+            from llama_cpp import Llama
+            
+            # Check if model exists locally first
+            model_filename = "Yi-1.5-9B-Chat-Q6_K.gguf"
+            local_model_path = os.path.join("./models", model_filename)
+            
+            if not os.path.exists(local_model_path):
+                logger.info("Downloading LLM model from HuggingFace...")
+                model_path = hf_hub_download(
+                    repo_id="IncreasingLoss/FineTunedTranslation_Yi-1.5-9B-Chat-Q6_K",
+                    filename=model_filename,
+                    local_dir="./models",
+                    local_dir_use_symlinks=False
+                )
+            else:
+                model_path = local_model_path
+                logger.info(f"Using existing model at: {model_path}")
+            
+            # Initialize with conservative settings for HF Spaces
+            llm = Llama(
+                model_path=model_path,
+                n_ctx=2048,  # Reduced context size
+                n_gpu_layers=0,  # CPU only for HF Spaces compatibility
+                verbose=False,
+                n_threads=2,  # Limit threads for shared environment
+                use_mmap=True,  # Memory mapping for efficiency
+                use_mlock=False  # Don't lock memory
+            )
+            logger.info("LLM model loaded successfully!")
+            
+        except Exception as e:
+            logger.error(f"Failed to load LLM: {e}")
+            llm = None
+    
+    return llm
 
 """yolo model"""
+logger.info("Loading YOLO model...")
 user_device = "cuda" if torch.cuda.is_available() else "cpu"
-detection_model = YOLO("model/yolo_chinese_m.pt").to(user_device).eval()
-
+logger.info(f"Using device: {user_device}")
 
+try:
+    detection_model = YOLO("model/yolo_chinese_m.pt").to(user_device).eval()
+    logger.info("YOLO model loaded successfully!")
+except Exception as e:
+    logger.error(f"Failed to load YOLO model: {e}")
+    raise
 
 """LW-Vit Classifier"""
 
@@ -195,14 +249,26 @@ class LW_ViT(nn.Module):
         x = self.classifier(x)
         return x
 
-classifier = LW_ViT()
-classifier.load_state_dict(torch.load("model/chinese_Character_Classification_handwritten_LW_ViT.pth", weights_only=True, map_location=user_device))
-classifier = classifier.to(user_device)
-classifier.eval()
+logger.info("Loading classifier model...")
+try:
+    classifier = LW_ViT()
+    classifier.load_state_dict(torch.load("model/chinese_Character_Classification_handwritten_LW_ViT.pth", weights_only=True, map_location=user_device))
+    classifier = classifier.to(user_device)
+    classifier.eval()
+    logger.info("Classifier model loaded successfully!")
+except Exception as e:
+    logger.error(f"Failed to load classifier: {e}")
+    raise
 
 """load classes dict"""
-with open("model/chinese_classes.json", "r", encoding="utf-8") as f:
-    classes_dict = json.load(f)
+logger.info("Loading classes dictionary...")
+try:
+    with open("model/chinese_classes.json", "r", encoding="utf-8") as f:
+        classes_dict = json.load(f)
+    logger.info(f"Loaded {len(classes_dict)} classes")
+except Exception as e:
+    logger.error(f"Failed to load classes dictionary: {e}")
+    raise
 
 """transforms"""
 class ToSquare(object):
@@ -230,7 +296,7 @@ class ToSquare(object):
         pad_w_extra = (max_dim - w) % 2
         pad_h_extra = (max_dim - h) % 2
         
-        # Create padding list (left, top, right, bottom) - FIX #1: Use list instead of tuple
+        # Create padding list (left, top, right, bottom)
         padding = [pad_w, pad_h, pad_w + pad_w_extra, pad_h + pad_h_extra]
         
         # Create new padded image
@@ -262,27 +328,9 @@ test_transforms = transforms.Compose([
     transforms.Normalize(mean=[0.5], std=[0.5])
 ])
 
-
-"""import llm from huggingface"""
-model_path = hf_hub_download(
-    repo_id="IncreasingLoss/FineTunedTranslation_Yi-1.5-9B-Chat-Q6_K",
-    filename="Yi-1.5-9B-Chat-Q6_K.gguf",
-    local_dir="./models",  # Optional: specify where to save
-    local_dir_use_symlinks=False
-)
-
-print(f"Model downloaded to: {model_path}")
-
-llm = Llama(
-    model_path=model_path,
-    n_ctx=4096,
-    n_gpu_layers=-1,  # Use all GPU layers
-    verbose=False
-)
-
 """app functions"""
 def convert_to_pil(inputimg):
-        # convert input to pil image
+    # convert input to pil image
     if inputimg is None:
         return None
     if isinstance(inputimg, np.ndarray):
@@ -304,8 +352,6 @@ def draw_bboxes(img, results):
         draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=2)
     return draw_img
 
-
-
 def crop_bboxes_by_reading_order(img, bboxes, written_vertical=False):
     """
     Crop and return image regions in proper reading order.
@@ -396,10 +442,8 @@ def crop_bboxes_by_reading_order(img, bboxes, written_vertical=False):
     
     return crops
 
-
 def move_slider(threshhold_slider, input_img):
     pil_img = convert_to_pil(input_img)
-    # FIX #2: Add null check for pil_img
     if pil_img is None:
         return None
     with torch.inference_mode():
@@ -411,163 +455,169 @@ def select_image(evt: gr.SelectData):
     selected_index = evt.index
     return example_images[selected_index]
 
-def generate_translation(prompt, max_new_tokens=512):
-    """Generate translation using Hugging Face model"""
-    
-    # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096-max_new_tokens)
-    
-    # Move to same device as model
-    inputs = {k: v.to(hf_model.device) for k, v in inputs.items()}
-    
-    # Generate response
-    with torch.no_grad():
-        outputs = hf_model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.3,
-            do_sample=True,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            stop_strings=["</s>", "\n\n", "Chinese:", "Chinese text:"],
-        )
-    
-    # Decode only the new tokens (exclude the input prompt)
-    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-    return response.strip()
-
-
 def translate_text(threshhold_slider, is_vertical, input_img):
     # convert input to pil image
     pil_img = convert_to_pil(input_img)
-    # FIX #2: Add null check for pil_img
     if pil_img is None:
-        return "", "", ""
+        return "", "", "Please upload an image first."
 
-    with torch.inference_mode():
-        results = detection_model(source=pil_img, conf=threshhold_slider)
-        sorted_cropped_images = crop_bboxes_by_reading_order(img=pil_img, bboxes=results[0].boxes.xyxy, written_vertical=is_vertical)
-        chinese_text = ""
-        for crop in sorted_cropped_images:
-            crop = test_transforms(crop)
-            crop = crop.to(user_device)  # Move tensor to device
-            crop_batch = crop.unsqueeze(dim=0)  # Add batch dimension
-
-            classifier.to(user_device)
+    try:
+        with torch.inference_mode():
+            results = detection_model(source=pil_img, conf=threshhold_slider)
+            sorted_cropped_images = crop_bboxes_by_reading_order(img=pil_img, bboxes=results[0].boxes.xyxy, written_vertical=is_vertical)
+            chinese_text = ""
+            for crop in sorted_cropped_images:
+                crop = test_transforms(crop)
+                crop = crop.to(user_device)  # Move tensor to device
+                crop_batch = crop.unsqueeze(dim=0)  # Add batch dimension
+
+                classifier.to(user_device)
+                
+                logits = classifier(crop_batch)
+                class_idx = logits.argmax(dim=1)
+                
+                class_idx_cpu = class_idx.cpu().item()  # Convert to CPU integer
+                chinese_charcter = classes_dict[class_idx_cpu]  # Use integer index for list
+                chinese_text += chinese_charcter
+
+        # Generate pinyin
+        pinyin_text = pinyin(chinese_text)
+        pinyin_sentence = " "
+        for pin in pinyin_text:
+            pinyin_sentence += f"{pin[0]} "
+
+        # Load LLM lazily only when needed for translation
+        current_llm = load_llm()
+        
+        if current_llm is None:
+            return chinese_text, pinyin_sentence, "Translation service unavailable - LLM failed to load."
+
+        # Generate translation
+        prompt = f""" You are a professional Chinese to English translator. 
+            1. Translate the following Chinese text to natural, fluent English:
+            "{chinese_text}"
+            2. Respond only with the translated English text.
+            English translation: """
+
+        try:
+            output = current_llm(
+                prompt, 
+                max_tokens=min(len(chinese_text)*3, 256),  # Conservative token limit
+                stop=["</s>", "\n\n", "Chinese:", "Chinese text:"],
+                echo=False,
+                temperature=0.3,
+                frequency_penalty=0.5,
+                presence_penalty=0.5,
+                top_p=0.9,
+                stream=False
+            )
             
-            logits = classifier(crop_batch)
-            class_idx = logits.argmax(dim=1)
+            # Extract the translation
+            translation = output["choices"][0]["text"].strip()
             
-            class_idx_cpu = class_idx.cpu().item()  # Convert to CPU integer
-            chinese_charcter = classes_dict[class_idx_cpu]  # Use integer index for list
-            chinese_text += chinese_charcter
-
-    pinyin_text = pinyin(chinese_text)
-    pinyin_sentence = " "
-    for pin in pinyin_text:
-        pinyin_sentence += f"{pin[0]} "
-
-
-    prompt = f""" You are a professional Chinese to English translator. 
-        1. Translate the following Chinese text to natural, fluent English:
-        "{chinese_text}"
-        2. Respond only with the translated English text.
-        English translation: """
-
-    output = llm(
-        prompt, 
-        max_tokens=min(len(chinese_text)*2, 512),  # More reasonable token limit with safety cap
-        stop=["</s>", "\n\n", "Chinese:", "Chinese text:"],  # Enhanced stop tokens
-        echo=False,
-        temperature=0.3,  # Slightly higher temperature for more natural language
-        frequency_penalty=.5,  # Penalize repetition of tokens - higher values discourage repetition
-        presence_penalty=.5,  # Penalize tokens that have appeared at all - prevents topic repetition
-        top_p=0.9,  # Nucleus sampling - restricts to top 90% most likely tokens for more coherent output
-        stream=False  # Explicitly disable streaming to get a dict response
-    )
-    
-    # Extract the translation
-    translation = output["choices"][0]["text"].strip()
-    try:
-        start = translation.index('"') + 1
-        # find index of last double-quote
-        end = translation.rindex('"')
-        translation = translation[start:end]    
-    except:
-        pass    
-
-    return chinese_text, pinyin_sentence, translation
-
+            # Clean up translation if it contains quotes
+            try:
+                if '"' in translation:
+                    start = translation.index('"') + 1
+                    end = translation.rindex('"')
+                    translation = translation[start:end]    
+            except:
+                pass    
+
+            return chinese_text, pinyin_sentence, translation
+            
+        except Exception as e:
+            logger.error(f"Translation failed: {e}")
+            return chinese_text, pinyin_sentence, f"Translation failed: {str(e)}"
+            
+    except Exception as e:
+        logger.error(f"OCR processing failed: {e}")
+        return "", "", f"OCR processing failed: {str(e)}"
 
 css = """
 .centered-examples {
     margin: 0 auto !important;
     justify-content: center !important;
     gap: 8px !important;
-    min-height: 150px !important;  /* Added minimum height */
+    min-height: 150px !important;
 }
 .centered-examples .thumb {
     height: 100px !important;
     width: 100px !important;
     object-fit: cover !important;
-    margin: 5px !important;  /* Added margin between thumbs */
+    margin: 5px !important;
 }
-/* 1) Global override: remove any forced sizing on .fixed-height anywhere */
 .fixed-height.svelte-842rpi.svelte-842rpi {
-  min-height: 0        !important;  /* cancel the global min-height */ 
-  max-height: none     !important;  /* cancel the global max-height */ 
-  height: auto         !important;  /* allow auto height */ 
+  min-height: 0        !important;
+  max-height: none     !important;
+  height: auto         !important;
 }
-/* 2) Same-query override: mirror Gradio's media query exactly */
 @media (min-width: 1280px) {
-  /* target the exact same class chain inside the breakpoint */
   .fixed-height.svelte-842rpi.svelte-842rpi {
-    min-height: 0      !important;  /* zero-out the 55vh/min-height there */ 
-    max-height: none   !important;  /* remove the viewport-height cap */ 
-    height: auto       !important;  /* let content dictate height */ 
+    min-height: 0      !important;
+    max-height: none   !important;
+    height: auto       !important;
   }
 }
 """
 
-
 """gradio app"""
+logger.info("Setting up Gradio interface...")
+
+# Check if examples directory exists
 example_dir = "examples"
-example_images = [os.path.join(example_dir, f) for f in os.listdir(example_dir) 
-                 if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
-with gr.Blocks(css=css) as program:
-    gr.Markdown("## DeepTranslate: Chinese OCR with translation to english")
+example_images = []
+if os.path.exists(example_dir):
+    example_images = [os.path.join(example_dir, f) for f in os.listdir(example_dir) 
+                     if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    logger.info(f"Found {len(example_images)} example images")
+else:
+    logger.warning(f"Examples directory '{example_dir}' not found")
+
+with gr.Blocks(css=css, title="DeepTranslate: Chinese OCR") as program:
+    gr.Markdown("## DeepTranslate: Chinese OCR with translation to English")
     gr.Markdown("Upload or select an image and move the slider to detect characters.")
     gr.Markdown("Make sure that the input image is high resolution and not rotated in any way!")
+    
+    # Show loading status
+    with gr.Row():
+        gr.Markdown("🟢 YOLO Model: Loaded | 🟢 Classifier: Loaded | 🟡 Translation: Loading on demand")
+    
     #inputs
     with gr.Column(scale=1):
-        gallery = gr.Gallery(value=example_images,             
-                            label="Example Images (Click to Select)",
-                            columns=6,
-                            allow_preview=False,
-                            elem_id="my_media_gallery",
-                            elem_classes=["centered-examples"])
+        if example_images:
+            gallery = gr.Gallery(value=example_images,             
+                                label="Example Images (Click to Select)",
+                                columns=6,
+                                allow_preview=False,
+                                elem_id="my_media_gallery",
+                                elem_classes=["centered-examples"])
         with gr.Row(scale=2):
             input_img = gr.Image(label="Input Image ")
-            detection_img = gr.Image(label="Detection Image", interactive= False) 
+            detection_img = gr.Image(label="Detection Image", interactive=False) 
 
     # slider and button 
     with gr.Column(scale=1):
         with gr.Row(scale=3):
-            threshhold_slider = gr.Slider(value=0.25, minimum=0, maximum=0.75, label="Detection Threshhold", step = 0.01)
+            threshhold_slider = gr.Slider(value=0.25, minimum=0, maximum=0.75, label="Detection Threshold", step=0.01)
             translate_button = gr.Button("Translate To English", variant="primary")
-            is_vertical = gr.Checkbox(value=False, label="Vertiacl Chinese Text?", interactive=True)
+            is_vertical = gr.Checkbox(value=False, label="Vertical Chinese Text?", interactive=True)
+    
     # outputs 
     with gr.Column(scale=1):
         with gr.Row(scale=3):
-            chinese_text = gr.TextArea(label="Chinese Text", max_lines= 1000, interactive=False)
-            pinyin_text  = gr.TextArea(label="Chinese Pinyin", max_lines= 1000, interactive=False)
-            english_text = gr.TextArea(label="English Text", max_lines= 1000, interactive=False)
-    #function, calling
-    threshhold_slider.change(fn= move_slider, inputs=[threshhold_slider, input_img], outputs=[detection_img])
-    translate_button.click(fn=translate_text, inputs= [threshhold_slider, is_vertical, input_img], outputs=[chinese_text, pinyin_text, english_text])
-    gallery.select(fn=select_image, inputs=None, outputs=input_img)
+            chinese_text = gr.TextArea(label="Chinese Text", max_lines=1000, interactive=False)
+            pinyin_text = gr.TextArea(label="Chinese Pinyin", max_lines=1000, interactive=False)
+            english_text = gr.TextArea(label="English Text", max_lines=1000, interactive=False)
+    
+    # function calling
+    threshhold_slider.change(fn=move_slider, inputs=[threshhold_slider, input_img], outputs=[detection_img])
+    translate_button.click(fn=translate_text, inputs=[threshhold_slider, is_vertical, input_img], outputs=[chinese_text, pinyin_text, english_text])
+    
+    if example_images:
+        gallery.select(fn=select_image, inputs=None, outputs=input_img)
+
+logger.info("Gradio interface ready!")
 
 if __name__ == "__main__":
-    program.launch(share=False)
+    program.launch(share=False, server_name="0.0.0.0", show_error=True)