Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download, snapshot_download
 import os
 import json
 import logging
@@ -11,23 +12,40 @@ logger = logging.getLogger(__name__)
 
 # Define model and checkpoint paths
 MODEL_REPO = "microsoft/CADFusion"
-CHECKPOINT_REVISION = "main"  # Use 'main' branch, which contains v1_1 checkpoint
+CHECKPOINT_REVISION = "main"
 CHECKPOINT_SUBFOLDER = "exp/model_ckpt/v1_1"
+LOCAL_CHECKPOINT_DIR = "./model_ckpt/v1_1"
+
+# Ensure local checkpoint directory exists
+os.makedirs(LOCAL_CHECKPOINT_DIR, exist_ok=True)
+
+# Download checkpoint files
+try:
+    logger.info("Downloading checkpoint files...")
+    snapshot_download(
+        repo_id=MODEL_REPO,
+        revision=CHECKPOINT_REVISION,
+        allow_patterns=f"{CHECKPOINT_SUBFOLDER}/*",
+        local_dir=LOCAL_CHECKPOINT_DIR,
+        local_dir_use_symlinks=False
+    )
+    logger.info("Checkpoint files downloaded successfully.")
+except Exception as e:
+    logger.error(f"Error downloading checkpoint files: {str(e)}")
+    raise e
 
 # Load model and tokenizer
 try:
     logger.info("Loading tokenizer...")
+    # Fallback to base Llama-3-8B tokenizer if CADFusion-specific config is missing
     tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_REPO,
-        revision=CHECKPOINT_REVISION,
-        subfolder=CHECKPOINT_SUBFOLDER,
+        "meta-llama/Meta-Llama-3-8B",
         trust_remote_code=True
     )
     logger.info("Loading model...")
+    # Attempt to load model from local checkpoint
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_REPO,
-        revision=CHECKPOINT_REVISION,
-        subfolder=CHECKPOINT_SUBFOLDER,
+        LOCAL_CHECKPOINT_DIR,
         torch_dtype=torch.float16,
         device_map="auto",
         trust_remote_code=True
@@ -35,7 +53,15 @@ try:
     logger.info("Model and tokenizer loaded successfully.")
 except Exception as e:
     logger.error(f"Error loading model or tokenizer: {str(e)}")
-    raise e
+    # Fallback to base Llama-3-8B model if local checkpoint fails
+    logger.info("Falling back to base Meta-Llama-3-8B model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Meta-Llama-3-8B",
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    logger.info("Fallback model loaded successfully.")
 
 # Function to generate CAD model from text description
 def generate_cad_model(text_description):
@@ -100,8 +126,8 @@ def create_gradio_interface():
         gr.Markdown("""
         **Note**: 
         - CADFusion is for research purposes only. Generated models may not be technically accurate and require validation.
-        - Ensure descriptions are clear and specific for best results.
-        - For more details, visit the [CADFusion GitHub repo](https://github.com/microsoft/CADFusion).
+        - This deployment uses a fallback to Meta-Llama-3-8B due to potential issues with the v1_1 checkpoint.
+        - For full functionality, refer to the [CADFusion GitHub repo](https://github.com/microsoft/CADFusion) for custom setup instructions.
         """)
     
     return demo