finalasing

app.py

@@ -6,6 +6,7 @@ import uvicorn
 import pika
 from fastapi import FastAPI, Body, Header, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from worker import main as rabbit_worker
 
 app = FastAPI()
 
@@ -55,4 +56,5 @@ async def process_pdf(
     }
 
 if __name__ == "__main__":
+    rabbit_worker()
     uvicorn.run(app, host="0.0.0.0", port=8000)

download_models_hf.py

@@ -60,6 +60,7 @@ if __name__ == '__main__':
     json_mods = {
         'models-dir': model_dir,
         'layoutreader-model-dir': layoutreader_model_dir,
+        'device': 'cuda'
     }
 
     download_and_modify_json(json_url, config_file, json_mods)

worker.py

@@ -122,7 +122,7 @@ def main():
     """
     Main entry: starts multiple worker threads to consume from the queue.
     """
-    num_workers = multiprocessing.cpu_count()
+    num_workers = 2 # hard code for now
     print(f"Starting {num_workers} workers")
 
     with ThreadPoolExecutor(max_workers=num_workers) as executor: