edit

app.py

@@ -79,18 +79,22 @@ def ocr_process(image):
         )
 
         # The processor outputs bfloat16 tensors, but model is float32.
-        # Convert all floating-point input tensors to float32.
-        for key in inputs:
-            if isinstance(inputs[key], torch.Tensor):
-                if inputs[key].is_floating_point():
-                    inputs[key] = inputs[key].float()
-
-        inputs = inputs.to("cpu")
+        # BatchFeature doesn't support in-place modification well,
+        # so rebuild as a plain dict with float32 tensors.
+        clean_inputs = {}
+        for k, v in inputs.items():
+            if isinstance(v, torch.Tensor):
+                if v.dtype == torch.bfloat16:
+                    clean_inputs[k] = v.to(torch.float32)
+                else:
+                    clean_inputs[k] = v
+            else:
+                clean_inputs[k] = v
 
         with torch.no_grad():
-            generated_ids = model.generate(**inputs, max_new_tokens=16384, do_sample=False)
+            generated_ids = model.generate(**clean_inputs, max_new_tokens=16384, do_sample=False)
 
-        input_ids = inputs["input_ids"]
+        input_ids = clean_inputs["input_ids"]
         generated_ids_trimmed = [
             out_ids[len(in_ids):] for in_ids, out_ids in zip(input_ids, generated_ids)
         ]
@@ -124,4 +128,4 @@ with gr.Blocks(title="HunyuanOCR") as demo:
     image_input.change(ocr_process, image_input, ocr_output)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0")
+    demo.launch(server_name="0.0.0.0")