Robust: Bypass PEFT generate wrapper to fix inference crash

app.py

@@ -15,6 +15,7 @@ IS_SPACE = "SPACE_ID" in os.environ
 print(f"System: Loading model... (Env: {'Hugging Face Space' if IS_SPACE else 'Local'})")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Load Processor
 try:
     processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
 except Exception:
@@ -22,7 +23,9 @@ except Exception:
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
     processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer)
 
+# Load Model
 base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID)
+# The PeftModel wrapper injects weights into base_model
 model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
 model.to(device)
 model.eval()
@@ -35,13 +38,20 @@ def predict(image):
         image = image.convert("RGB")
         pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
         
-        # --- VISION-SPECIFIC GENERATE CALL ---
-        # Vision models like TrOCR require the 'pixel_values' keyword explicitly
+        # --- THE ROBUST FIX ---
+        # We call .base_model.generate() directly.
+        # This bypasses the buggy PEFT wrapper while still using the LoRA weights.
+        # We also add max_new_tokens for a better result.
         with torch.no_grad():
-            generated_ids = model.generate(pixel_values=pixel_values)
+            generated_ids = model.base_model.generate(
+                pixel_values=pixel_values,
+                max_new_tokens=64
+            )
             generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return generated_text
     except Exception as e:
+        import traceback
+        print(traceback.format_exc()) # Log the full error to the Space logs
         return f"Error during inference: {str(e)}"
 
 # Interface setup