From bfa6c0e1c1aac7485b938effde9a85112a67ed80 Mon Sep 17 00:00:00 2001
From: julian <jcloiacon@gmail.com>
Date: Tue, 28 Jan 2025 12:49:39 -0500
Subject: [PATCH] Adds MODEL_NAME as an Environment variable for Docker
 context. Defaults to the same 1.3B model

---
 demo/fastapi_app.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/demo/fastapi_app.py b/demo/fastapi_app.py
index 54d4a22..f98579b 100644
--- a/demo/fastapi_app.py
+++ b/demo/fastapi_app.py
@@ -6,11 +6,12 @@ from janus.models import MultiModalityCausalLM, VLChatProcessor
 from PIL import Image
 import numpy as np
 import io
+import os
 
 app = FastAPI()
 
 # Load model and processor
-model_path = "deepseek-ai/Janus-1.3B"
+model_path = os.getenv("MODEL_NAME", "deepseek-ai/Janus-1.3B")
 config = AutoConfig.from_pretrained(model_path)
 language_config = config.language_config
 language_config._attn_implementation = 'eager'