Update app.py

app.py

@@ -5,6 +5,7 @@ import jinja2
 import pdfkit
 import torch
 import logging
+import subprocess
 from threading import Thread
 from flask import Flask, request, send_file, jsonify
 from flask_cors import CORS
@@ -29,10 +30,18 @@ model_loaded = False
 load_error = None
 generator = None
 
+# Find wkhtmltopdf path
+WKHTMLTOPDF_PATH = '/usr/bin/wkhtmltopdf'
+if not os.path.exists(WKHTMLTOPDF_PATH):
+    # Try to find it using which
+    try:
+        WKHTMLTOPDF_PATH = subprocess.check_output(['which', 'wkhtmltopdf']).decode().strip()
+    except:
+        app.logger.warning("Could not find wkhtmltopdf path. Using default.")
+        WKHTMLTOPDF_PATH = 'wkhtmltopdf'
+
 # Configure wkhtmltopdf
-# Use xvfb-run for headless PDF generation
-WKHTMLTOPDF_CMD = 'xvfb-run -a wkhtmltopdf'
-pdf_config = pdfkit.configuration(wkhtmltopdf=WKHTMLTOPDF_CMD)
+pdf_config = pdfkit.configuration(wkhtmltopdf=WKHTMLTOPDF_PATH)
 
 def load_model():
     global model_loaded, load_error, generator
@@ -217,9 +226,27 @@ def generate_pdf():
         try:
             with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as f:
                 pdf_path = f.name
+            
+            # Generate PDF using xvfb-run as a separate process
+            html_path = pdf_path + '.html'
+            with open(html_path, 'w', encoding='utf-8') as f:
+                f.write(html)
                 
-            # Generate PDF using wkhtmltopdf with xvfb
-            pdfkit.from_string(html, pdf_path, options=options, configuration=pdf_config)
+            command = ['xvfb-run', '-a', WKHTMLTOPDF_PATH] + \
+                      [f'--{k}={v}' for k, v in options.items() if v] + \
+                      [html_path, pdf_path]
+            
+            app.logger.info(f"Running command: {' '.join(command)}")
+            result = subprocess.run(command, capture_output=True, text=True)
+            
+            if result.returncode != 0:
+                app.logger.error(f"PDF generation command failed: {result.stderr}")
+                # Fallback to direct pdfkit if available
+                app.logger.info("Trying fallback PDF generation with pdfkit")
+                pdfkit.from_string(html, pdf_path, options=options, configuration=pdf_config)
+            
+            # Clean up HTML file
+            os.remove(html_path)
             
             app.logger.info(f"PDF generated successfully at {pdf_path}")
             return send_file(pdf_path, mimetype='application/pdf', as_attachment=True,