Adds MODEL_NAME as an Environment variable for Docker context. Defaults to the same 1.3B model

This commit is contained in:
julian 2025-01-28 12:49:39 -05:00
parent e471403e44
commit bfa6c0e1c1

View File

@ -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'