refactor: update inference to use priming by default

main.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import logging 
 import time 
 from contextlib import asynccontextmanager
-from inference_utils import construct_alphabet_list, convert_offsets_to_absolute_coords, encode_text, get_alphabet_map
+from inference_utils import PrimingData, construct_alphabet_list, convert_offsets_to_absolute_coords, encode_text, get_alphabet_map, load_priming_data
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__) 
@@ -18,13 +18,13 @@ QUANTIZED_MODEL_NAME = "model.scripted.quantized.pt"
 SCRIPTED_MODEL_NAME = "model.scripted.pt" 
 METADATA_MODEL_NAME = "model.pt" 
 
-scripted_model: Optional[torch.jit.ScriptModule] = None 
-model_metadata: Optional[dict] = None 
-device: Optional[torch.device] = None 
-alphabet_map: Optional[dict[str, int]] = None 
+SCRIPTED_MODEL: Optional[torch.jit.ScriptModule] = None 
+MODEL_METADATA: Optional[dict] = None 
+DEVICE: Optional[torch.device] = None 
+ALPHABET_MAP: Optional[dict[str, int]] = None 
 ALPHABET_LIST: Optional[list[str]] = None 
 ALPHABET_SIZE: Optional[int] = None 
-max_text_len: Optional[int] = None 
+MAX_TEXT_LEN: Optional[int] = None 
 output_mixture_components: Optional[int] = None # To store num_mixtures for GMM sampling
 lstm_size: Optional[int] = None 
 attention_mixture_components: Optional[int] = None 
@@ -55,15 +55,15 @@ class HealthResponse(BaseModel):
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Lifespan context manager for startup and shutdown events"""
-    global scripted_model, model_metadata, device, alphabet_map, max_text_len, ALPHABET_LIST, output_mixture_components, lstm_size, attention_mixture_components, ALPHABET_SIZE
+    global SCRIPTED_MODEL, MODEL_METADATA, DEVICE, ALPHABET_MAP, MAX_TEXT_LEN, ALPHABET_LIST, output_mixture_components, lstm_size, attention_mixture_components, ALPHABET_SIZE
     logger.info("Attempting to load model resources during startup") 
     try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        logger.info(f"Using device: {device}")
+        DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {DEVICE}")
         
         scripted_model_path = MODEL_DIR / SCRIPTED_MODEL_NAME
         metadata_model_path = MODEL_DIR / METADATA_MODEL_NAME
-        if device.type == "cpu":
+        if DEVICE.type == "cpu":
             scripted_model_path = MODEL_DIR / QUANTIZED_MODEL_NAME
         
         if  not scripted_model_path.exists():
@@ -74,18 +74,18 @@ async def lifespan(app: FastAPI):
             raise FileNotFoundError(f"Metadata model file not found at {metadata_model_path}")
 
         # Load the traced model
-        scripted_model = torch.jit.load(scripted_model_path, map_location=device)
-        if scripted_model:
-            scripted_model.eval()
+        SCRIPTED_MODEL = torch.jit.load(scripted_model_path, map_location=DEVICE)
+        if SCRIPTED_MODEL:
+            SCRIPTED_MODEL.eval()
             logger.info(f"Traced model loaded successfully from {scripted_model_path}")
 
         # Load the metadata
-        model_metadata = torch.load(metadata_model_path, map_location='cpu')
-        if model_metadata:
+        MODEL_METADATA = torch.load(metadata_model_path, map_location='cpu')
+        if MODEL_METADATA:
             logger.info(f"Model metadata loaded successfully from {metadata_model_path}")
-            logger.info(f"Model metadata keys: {list(model_metadata.keys())}")
+            logger.info(f"Model metadata keys: {list(MODEL_METADATA.keys())}")
             
-            config_full = model_metadata['config_full']
+            config_full = MODEL_METADATA['config_full']
             if not config_full or not isinstance(config_full, dict):
                 raise ValueError(f"Key `config_full` not found or not a dict")
             
@@ -95,7 +95,7 @@ async def lifespan(app: FastAPI):
             if not dataset_config or not isinstance(dataset_config, dict):
                 raise ValueError(f"Key `dataset` not found or not a dict in config_full")
             alphabet_str = dataset_config['alphabet_string']
-            max_text_len = dataset_config['max_text_len']
+            MAX_TEXT_LEN = dataset_config['max_text_len']
             output_mixture_components = model_params['output_mixture_components']
 
             lstm_size = model_params['lstm_size']
@@ -103,7 +103,7 @@ async def lifespan(app: FastAPI):
 
             ALPHABET_LIST = construct_alphabet_list(alphabet_str)
             ALPHABET_SIZE = len(ALPHABET_LIST)
-            alphabet_map = get_alphabet_map(ALPHABET_LIST) 
+            ALPHABET_MAP = get_alphabet_map(ALPHABET_LIST) 
             
             logger.info(f"Alphabet created. Size: {len(ALPHABET_LIST)}")
             logger.info("Model resources are loaded and ready")
@@ -112,16 +112,16 @@ async def lifespan(app: FastAPI):
 
     except Exception as e:
         logger.error(f"Error loading model resources: {e}", exc_info=True)
-        scripted_model = None
-        model_metadata = None
+        SCRIPTED_MODEL = None
+        MODEL_METADATA = None
         raise
     
     yield
     
     # Cleanup on shutdown
     logger.info("Shutting down API and cleaning up resources")
-    scripted_model = None 
-    model_metadata = None 
+    SCRIPTED_MODEL = None 
+    MODEL_METADATA = None 
 
 app = FastAPI(
     title="Scriptify API", 
@@ -145,32 +145,31 @@ async def read_root():
 
 @app.get("/health", response_model=HealthResponse, tags=["General"])
 async def health_check():
-    global scripted_model, model_metadata, device, alphabet_map, max_text_len, ALPHABET_LIST 
+    global SCRIPTED_MODEL, MODEL_METADATA, DEVICE, ALPHABET_MAP, MAX_TEXT_LEN, ALPHABET_LIST 
     
-    is_healthy = all([scripted_model, model_metadata, device, alphabet_map, max_text_len, ALPHABET_LIST])
+    is_healthy = all([SCRIPTED_MODEL, MODEL_METADATA, DEVICE, ALPHABET_MAP, MAX_TEXT_LEN, ALPHABET_LIST])
     
     return HealthResponse(
         status="healthy" if is_healthy else "unhealthy",
-        model_loaded=bool(scripted_model),
-        device=str(device) if device else "unknown",
-        model_metadata_keys=list(model_metadata.keys()) if model_metadata else None,
+        model_loaded=bool(SCRIPTED_MODEL),
+        device=str(DEVICE) if DEVICE else "unknown",
+        model_metadata_keys=list(MODEL_METADATA.keys()) if MODEL_METADATA else None,
     )
 
-def text_to_tensor(text: str, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+def text_to_tensor(text: str, max_text_length: int, add_eos: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
     """Convert text to tensor format expected by the model"""
-    global alphabet_map, max_text_len  
-    if alphabet_map is None:
+    if ALPHABET_MAP is None:
         raise ValueError("Alphabet map not initialized during api startup")
-    if max_text_len is None:
-        raise ValueError("`max_text_len` is not initialized during api startup")
+
     padded_encoded_np, true_length = encode_text(
         text=text, 
-        char_to_index_map=alphabet_map, 
-        max_length=max_text_len 
+        char_to_index_map=ALPHABET_MAP, 
+        max_length=max_text_length, 
+        add_eos = add_eos 
     )
 
-    char_seq = torch.from_numpy(padded_encoded_np).to(device=device, dtype=torch.long)
-    char_len = torch.tensor([true_length], device=device, dtype=torch.long)
+    char_seq = torch.from_numpy(padded_encoded_np).to(device=DEVICE, dtype=torch.long)
+    char_len = torch.tensor([true_length], device=DEVICE, dtype=torch.long)
 
     return char_seq, char_len
 
@@ -179,20 +178,38 @@ def generate_strokes(
     char_lengths: torch.Tensor,
     max_gen_len: int,
     api_bias: float,
-    current_device: torch.device
+    style: Optional[int] = None
 ) -> list[list[float]]:
     """Generate strokes using the model's built-in sample method"""
-    global scripted_model
-    if scripted_model is None:
+    global SCRIPTED_MODEL
+    if SCRIPTED_MODEL is None:
         raise ValueError("Scripted model not initialized.")
     
+    primingData = None 
+
+    if style is not None: 
+        priming_text, priming_strokes = load_priming_data(style) 
+
+        priming_text_tensor, priming_text_len_tensor = text_to_tensor(
+            priming_text, max_text_length=len(priming_text), add_eos=False)
+
+        priming_stroke_tensor = torch.tensor(priming_strokes, 
+                                             dtype=torch.float32, 
+                                             device=DEVICE).unsqueeze(dim=0)
+        
+        primingData = PrimingData(priming_stroke_tensor,
+                                  char_seq_tensors=priming_text_tensor, 
+                                  char_seq_lengths=priming_text_len_tensor)
+        
+    
     with torch.no_grad():
         try:
-            stroke_tensors = scripted_model.sample(
+            stroke_tensors = SCRIPTED_MODEL.sample(
                 char_seq, 
                 char_lengths, 
                 max_length=max_gen_len, 
-                bias=api_bias
+                bias=api_bias,
+                prime=primingData
             )
             
             if len(stroke_tensors) == 1 and stroke_tensors[0].dim() == 2:
@@ -217,21 +234,24 @@ def generate_strokes(
             
 @app.post("/generate", response_model=HandwritingResponse, tags=["Generation"])
 async def generate_handwriting_endpoint(request: HandwritingRequest):
-    if not all([scripted_model, model_metadata, device, alphabet_map, max_text_len]):
+    if not all([SCRIPTED_MODEL, MODEL_METADATA, DEVICE, ALPHABET_MAP, MAX_TEXT_LEN]):
         logger.error("API not fully initialized. Check /health endpoint.")
         raise HTTPException(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail="Model or required resources not loaded."
         )
     
-    assert device is not None, "Device is None inside generate_handwriting"
+    assert DEVICE is not None, "Device is None inside generate_handwriting"
     start_time = time.time()
     
     try:
-        char_seq_tensor, char_lengths_tensor = text_to_tensor(request.text, device)
+        char_seq_tensor, char_lengths_tensor = text_to_tensor(request.text, max_text_length=MAX_TEXT_LEN) # type: ignore 
 
         relative_stroke_offsets = generate_strokes(
-            char_seq_tensor, char_lengths_tensor, request.max_length, request.bias, device
+            char_seq_tensor, char_lengths_tensor, 
+            request.max_length,
+            request.bias,
+            style=1 #TODO: style is hardcode since the current version is hosted on cpu 
         )
 
         if not relative_stroke_offsets: