remove task type extract text

app.py

@@ -14,7 +14,7 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
     use_safetensors=True,
 )
-model = model.eval().cuda().to(torch.bfloat16)
+model = model.eval()
 
 
 @spaces.GPU
@@ -27,6 +27,9 @@ def process_image(image, model_size, task_type):
         model_size: Model size configuration
         task_type: OCR task type
     """
+    # 在 GPU 函数内部移动模型到 GPU
+    model_gpu = model.cuda().to(torch.bfloat16)
+
     # Create temporary directory for output
     with tempfile.TemporaryDirectory() as output_path:
         # Set prompt based on task type
@@ -34,8 +37,6 @@ def process_image(image, model_size, task_type):
             prompt = "<image>\nFree OCR. "
         elif task_type == "Convert to Markdown":
             prompt = "<image>\n<|grounding|>Convert the document to markdown. "
-        elif task_type == "Extract Text":
-            prompt = "<image>\nExtract all text from the image. "
         else:
             prompt = "<image>\nFree OCR. "
 
@@ -59,7 +60,7 @@ def process_image(image, model_size, task_type):
         config = size_configs.get(model_size, size_configs["Gundam (Recommended)"])
 
         # Run inference
-        result = model.infer(
+        result = model_gpu.infer(
             tokenizer,
             prompt=prompt,
             image_file=temp_image_path,
@@ -71,6 +72,7 @@ def process_image(image, model_size, task_type):
             test_compress=True,
         )
 
+        print(f"====\nresult: {result}\n====\n")
         return result
 
 
@@ -105,7 +107,7 @@ with gr.Blocks(title="DeepSeek-OCR") as demo:
             )
 
             task_type = gr.Dropdown(
-                choices=["Free OCR", "Convert to Markdown", "Extract Text"],
+                choices=["Free OCR", "Convert to Markdown"],
                 value="Convert to Markdown",
                 label="Task Type",
             )