update: nougat gpu

app.py

@@ -87,43 +87,107 @@ except:
 if not terminators:
     terminators = [2]  # 使用常见的</s>标记ID作为默认值
 
-# 优化后的GPU-based Nougat PDF处理
+# 使用CUDA运行Nougat的PDF处理函数
+def process_pdf_with_nougat_gpu(pdf_path, output_dir=None):
+    """使用GPU运行Nougat处理PDF文件"""
+    try:
+        # 如果未指定输出目录，使用PDF所在目录
+        if output_dir is None:
+            output_dir = os.path.dirname(pdf_path)
+        
+        # 设置CUDA环境变量
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = "0"  # 使用第一个GPU
+        env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+        
+        # 执行带有GPU支持的Nougat命令
+        print(f"使用GPU运行Nougat: {pdf_path}")
+        cmd = ["nougat", pdf_path, "-o", output_dir, "--device", "cuda"]
+        
+        # 执行命令并捕获输出
+        result = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            env=env,
+            timeout=300  # 5分钟超时
+        )
+        
+        # 检查命令执行结果
+        if result.returncode != 0:
+            print(f"Nougat GPU处理失败: {result.stderr}")
+            return None, result.stderr
+        
+        # 获取生成的markdown文件路径
+        base_name = os.path.basename(pdf_path)
+        name_without_ext = os.path.splitext(base_name)[0]
+        markdown_path = os.path.join(output_dir, f"{name_without_ext}.mmd")
+        
+        # 检查markdown文件是否生成
+        if not os.path.exists(markdown_path):
+            return None, "Nougat处理完成，但未找到生成的Markdown文件"
+        
+        # 读取markdown内容
+        with open(markdown_path, "r", encoding="utf-8") as f:
+            markdown_content = f.read()
+        
+        return markdown_content, None
+    
+    except subprocess.TimeoutExpired:
+        return None, "Nougat处理超时"
+    
+    except Exception as e:
+        import traceback
+        error = f"Nougat处理异常: {str(e)}\n{traceback.format_exc()}"
+        print(error)
+        return None, error
+
+# 使用Python API的GPU处理方式
 @spaces.GPU(stateless=True)
-def process_pdf_with_nougat(pdf_path):
-    """使用Nougat处理PDF文件 (在GPU环境中运行)"""
+def process_pdf_with_nougat_api(pdf_path):
+    """使用Nougat Python API与GPU处理PDF文件"""
     try:
-        # 方法1: 使用Nougat Python API
-        try:
-            from nougat import NougatModel
-            from nougat.utils.checkpoint import get_checkpoint
-            from nougat.dataset.rasterize import rasterize_paper
-            import torch
-            
-            # 初始化模型并移至GPU
-            ckpt = get_checkpoint()
-            model = NougatModel.from_pretrained(ckpt)
-            device = torch.device("cuda")
-            model = model.to(device)
-            
-            # 处理PDF
-            markdown_content = ""
-            for page_idx, page in enumerate(rasterize_paper(pdf_path)):
-                page = page.to(device)
-                markdown = model.inference(page)
-                markdown_content += f"--- Page {page_idx+1} ---\n{markdown}\n\n"
-            
-            print("成功使用Nougat Python API处理PDF")
-            return markdown_content
-        except Exception as api_error:
-            print(f"Nougat Python API处理失败: {str(api_error)}")
-            raise api_error
-            
+        # 导入必要的库
+        from nougat import NougatModel
+        from nougat.utils.checkpoint import get_checkpoint
+        from nougat.dataset.rasterize import rasterize_paper
+        import torch
+        
+        # 确保GPU可用
+        if not torch.cuda.is_available():
+            return None, "GPU不可用，无法使用Nougat API处理PDF"
+        
+        # 显示GPU信息
+        device_count = torch.cuda.device_count()
+        device_name = torch.cuda.get_device_name(0) if device_count > 0 else "Unknown"
+        print(f"使用GPU: {device_name}, 可用GPU数量: {device_count}")
+        
+        # 初始化模型并移至GPU
+        ckpt = get_checkpoint()
+        model = NougatModel.from_pretrained(ckpt)
+        device = torch.device("cuda")
+        model = model.to(device)
+        
+        # 处理PDF
+        markdown_content = ""
+        pages = list(rasterize_paper(pdf_path))
+        
+        # 使用tqdm显示进度
+        for page_idx, page in enumerate(tqdm(pages, desc="处理PDF页面")):
+            page = page.to(device)
+            markdown = model.inference(page)
+            markdown_content += f"--- Page {page_idx+1} ---\n{markdown}\n\n"
+        
+        return markdown_content, None
+    
     except Exception as e:
         import traceback
-        print(f"GPU PDF处理失败: {str(e)}\n{traceback.format_exc()}")
-        raise e
+        error = f"Nougat API处理异常: {str(e)}\n{traceback.format_exc()}"
+        print(error)
+        return None, error
 
-# 添加优化后的PDF转换为Markdown函数
+# 添加PDF转换为Markdown函数
 def convert_pdf_to_markdown(pdf_file):
     """使用Nougat将PDF转换为Markdown (GPU优化版)"""
     if pdf_file is None:
@@ -141,73 +205,34 @@ def convert_pdf_to_markdown(pdf_file):
             with open(temp_pdf_path, "wb") as f:
                 f.write(pdf_file)
             
-            # 首先尝试使用GPU命令行方式处理PDF
-            output_dir = temp_dir
-            print(f"执行: nougat {temp_pdf_path} -o {output_dir}")
-            
-            # 设置GPU环境变量
-            env = os.environ.copy()
-            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+            # 方法1: 首先尝试使用命令行GPU方式
+            print("方法1: 尝试使用命令行GPU方式处理PDF...")
+            markdown_content, error = process_pdf_with_nougat_gpu(temp_pdf_path, temp_dir)
             
-            try:
-                result = subprocess.run(
-                    ["nougat", temp_pdf_path, "-o", output_dir], 
-                    stdout=subprocess.PIPE, 
-                    stderr=subprocess.PIPE,
-                    text=True,
-                    timeout=180,
-                    env=env
-                )
-                
-                # 检查命令行转换是否成功
-                if result.returncode == 0:
-                    # 读取生成的Markdown文件
-                    markdown_path = os.path.join(output_dir, "temp.mmd")
-                    if os.path.exists(markdown_path):
-                        with open(markdown_path, "r", encoding="utf-8") as f:
-                            markdown_content = f.read()
-                        
-                        # 限制文本长度
-                        if len(markdown_content) > 20000:
-                            markdown_content = markdown_content[:20000] + "\n\n...(Markdown内容已截断)"
-                        
-                        status = f"PDF已成功转换为Markdown (GPU命令行): 生成了{len(markdown_content)}个字符"
-                        return markdown_content, status
-                
-                # 如果命令行方式失败，尝试空间GPU API
-                print("命令行转换失败，正在尝试使用GPU API方式处理PDF...")
-                markdown_content = process_pdf_with_nougat(temp_pdf_path)
-                
-                # 限制文本长度
-                if len(markdown_content) > 20000:
-                    markdown_content = markdown_content[:20000] + "\n\n...(Markdown内容已截断)"
-                
-                status = f"PDF已成功转换为Markdown (GPU API): 生成了{len(markdown_content)}个字符"
-                return markdown_content, status
-                
-            except subprocess.TimeoutExpired:
-                print("命令行处理超时，尝试使用GPU API...")
-                # 尝试使用GPU API
-                markdown_content = process_pdf_with_nougat(temp_pdf_path)
-                
+            if markdown_content is not None:
                 # 限制文本长度
                 if len(markdown_content) > 20000:
                     markdown_content = markdown_content[:20000] + "\n\n...(Markdown内容已截断)"
                 
-                status = f"PDF已成功转换为Markdown (GPU API): 生成了{len(markdown_content)}个字符"
+                status = f"PDF已成功转换为Markdown (GPU命令行): 生成了{len(markdown_content)}个字符"
                 return markdown_content, status
-                
-            except Exception as cmd_error:
-                print(f"命令行处理失败: {str(cmd_error)}")
-                # 尝试使用GPU API
-                markdown_content = process_pdf_with_nougat(temp_pdf_path)
-                
+            
+            # 方法2: 如果命令行方式失败，尝试使用Python API方式
+            print(f"方法1失败: {error}")
+            print("方法2: 尝试使用Python API GPU方式处理PDF...")
+            
+            markdown_content, api_error = process_pdf_with_nougat_api(temp_pdf_path)
+            
+            if markdown_content is not None:
                 # 限制文本长度
                 if len(markdown_content) > 20000:
                     markdown_content = markdown_content[:20000] + "\n\n...(Markdown内容已截断)"
                 
                 status = f"PDF已成功转换为Markdown (GPU API): 生成了{len(markdown_content)}个字符"
                 return markdown_content, status
+            
+            # 所有方法都失败
+            return "", f"PDF转换失败: 所有GPU方法都失败了\n命令行错误: {error}\nAPI错误: {api_error}"
     
     except Exception as e:
         import traceback