Mitigate peak memory usage

app.py

@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 # Required for access to a gated model
 load_dotenv()
-hf_token = os.getenv("HUGGING_FACE_TOKEN", None)
+hf_token = os.getenv("HF_TOKEN", None)
 if hf_token is not None:
     login(token=hf_token)
 
@@ -53,8 +53,12 @@ async def lifespan(app: FastAPI):
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-        ).to(device)
-        logger.info("Model loaded successfully!")
+            device_map="auto",        # Automatically splits model across CPU/GPU
+            low_cpu_mem_usage=True,   # Avoids unnecessary CPU memory duplication
+            offload_folder="offload", # Temporary offload to disk
+        )
+        #.to(device)
+        logger.info(f"Model loaded successfully! ({device})")
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
         raise e