Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import snapshot_download
 import os
 import json
 import logging
@@ -15,6 +15,7 @@ MODEL_REPO = "microsoft/CADFusion"
 CHECKPOINT_REVISION = "main"
 CHECKPOINT_SUBFOLDER = "exp/model_ckpt/v1_1"
 LOCAL_CHECKPOINT_DIR = "./model_ckpt/v1_1"
+FALLBACK_MODEL = "meta-llama/Llama-2-7b"
 
 # Ensure local checkpoint directory exists
 os.makedirs(LOCAL_CHECKPOINT_DIR, exist_ok=True)
@@ -36,32 +37,37 @@ except Exception as e:
 
 # Load model and tokenizer
 try:
-    logger.info("Loading tokenizer...")
-    # Fallback to base Llama-3-8B tokenizer if CADFusion-specific config is missing
+    logger.info("Loading tokenizer from local checkpoint...")
     tokenizer = AutoTokenizer.from_pretrained(
-        "meta-llama/Meta-Llama-3-8B",
+        LOCAL_CHECKPOINT_DIR,
         trust_remote_code=True
     )
-    logger.info("Loading model...")
-    # Attempt to load model from local checkpoint
+    logger.info("Loading model from local checkpoint...")
     model = AutoModelForCausalLM.from_pretrained(
         LOCAL_CHECKPOINT_DIR,
         torch_dtype=torch.float16,
         device_map="auto",
         trust_remote_code=True
     )
-    logger.info("Model and tokenizer loaded successfully.")
+    logger.info("Model and tokenizer loaded successfully from local checkpoint.")
 except Exception as e:
-    logger.error(f"Error loading model or tokenizer: {str(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.")
+    logger.error(f"Error loading from local checkpoint: {str(e)}")
+    logger.info(f"Falling back to {FALLBACK_MODEL}...")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            FALLBACK_MODEL,
+            trust_remote_code=True
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            FALLBACK_MODEL,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        logger.info(f"Fallback model {FALLBACK_MODEL} loaded successfully.")
+    except Exception as fallback_e:
+        logger.error(f"Error loading fallback model: {str(fallback_e)}")
+        raise fallback_e
 
 # Function to generate CAD model from text description
 def generate_cad_model(text_description):
@@ -125,9 +131,10 @@ def create_gradio_interface():
         
         gr.Markdown("""
         **Note**: 
+        - This deployment may use a fallback model (Llama-2-7b) due to issues with the CADFusion v1_1 checkpoint.
         - CADFusion is for research purposes only. Generated models may not be technically accurate and require validation.
-        - 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.
+        - For full CADFusion functionality, follow the setup instructions in the [CADFusion GitHub repo](https://github.com/microsoft/CADFusion).
+        - Contact Shizhao Sun (shizsu@microsoft.com) for checkpoint access or issues.
         """)
     
     return demo