Update app.py

app.py

@@ -252,40 +252,32 @@ def get_hub_or_local_checkpoint(repo_id: str, filename: str):
     return ckpt_path
 
 def download_gemma_model(repo_id: str):
-    """
-    Download the whole repo into a normal app folder instead of ~/.cache/huggingface.
-    """
-    local_repo_dir = MODEL_DIR / repo_id.replace("/", "__")
-    local_repo_dir.mkdir(parents=True, exist_ok=True)
-
-    log_event(f"Downloading Gemma model from {repo_id} into {local_repo_dir}...")
-    local_dir = snapshot_download(
-        repo_id=repo_id,
-        local_dir=local_repo_dir,
-    )
-    log_event(f"Gemma model downloaded to {local_dir}")
+    """Download the full Gemma model directory."""
+    print(f"Downloading Gemma model from {repo_id}...")
+    local_dir = snapshot_download(repo_id=repo_id)
+    print(f"Gemma model downloaded to {local_dir}")
     return local_dir
 
 # Initialize model ledger and text encoder at startup (load once, keep in memory)
-log_event("=" * 80)
-log_event("Loading Gemma Text Encoder...")
-log_event("=" * 80)
+print("=" * 80)
+print("Loading Gemma Text Encoder...")
+print("=" * 80)
 
 checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)
 gemma_local_path = download_gemma_model(DEFAULT_GEMMA_REPO_ID)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda"
 
-log_event(f"Initializing text encoder with:")
-log_event(f"  checkpoint_path={checkpoint_path}")
-log_event(f"  gemma_root={gemma_local_path}")
-log_event(f"  device={device}")
+print(f"Initializing text encoder with:")
+print(f"  checkpoint_path={checkpoint_path}")
+print(f"  gemma_root={gemma_local_path}")
+print(f"  device={device}")
 
 
 model_ledger = ModelLedger(
     dtype=torch.bfloat16,
     device=device,
     checkpoint_path=checkpoint_path,
-    gemma_root_path=gemma_local_path,
+    gemma_root_path=DEFAULT_GEMMA_REPO_ID,
     local_files_only=False
 )