Update app.py

app.py

@@ -1,14 +1,21 @@
 import gradio as gr
 import torch
 from transformers import AutoProcessor, PaddleOCRVLForConditionalGeneration
+import time
 
 print("Loading model...")
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+
 model = PaddleOCRVLForConditionalGeneration.from_pretrained(
     "PaddlePaddle/PaddleOCR-VL",
     torch_dtype=torch.bfloat16,
     device_map="auto"
 )
 processor = AutoProcessor.from_pretrained("PaddlePaddle/PaddleOCR-VL")
+
+print(f"Model device: {model.device}")
 print("Model loaded!")
 
 PROMPTS = {
@@ -22,11 +29,17 @@ def predict(image, task):
     if image is None:
         return "Загрузите изображение"
     
+    print(f"[{time.strftime('%H:%M:%S')}] Processing started...")
+    print(f"Image size: {image.size}")
+    
+    t0 = time.time()
+    
     messages = [{"role": "user", "content": [
         {"type": "image", "image": image},
         {"type": "text", "text": PROMPTS[task]},
     ]}]
     
+    print(f"[{time.strftime('%H:%M:%S')}] Tokenizing...")
     inputs = processor.apply_chat_template(
         messages,
         tokenize=True,
@@ -35,16 +48,19 @@ def predict(image, task):
         return_tensors="pt"
     ).to(model.device)
     
-    generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    print(f"[{time.strftime('%H:%M:%S')}] Generating... (input shape: {inputs.input_ids.shape})")
+    
+    with torch.inference_mode():
+        generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    
+    print(f"[{time.strftime('%H:%M:%S')}] Decoding...")
     generated_ids_trimmed = [
         out_ids[len(in_ids):] 
         for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
     ]
-    result = processor.batch_decode(
-        generated_ids_trimmed, 
-        skip_special_tokens=True
-    )[0]
+    result = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True)[0]
     
+    print(f"[{time.strftime('%H:%M:%S')}] Done in {time.time()-t0:.1f}s")
     return result
 
 gr.Interface(