Update app.py

app.py

@@ -3,20 +3,30 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import warnings
 
-# Suppress the specific warning
+# Suppress warnings
 warnings.filterwarnings("ignore", category=FutureWarning, module="huggingface_hub")
+warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
 
-# Download and load the model
-model_name = "baidu/ERNIE-4.5-21B-A3B-Thinking"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
+try:
+    # Download and load the model
+    model_name = "baidu/ERNIE-4.5-21B-A3B-Thinking"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    model_loaded = True
+except Exception as e:
+    print(f"Model loading failed: {e}")
+    model_loaded = False
 
 def generate_code(prompt):
     """Generate HTML code using ERNIE model"""
+    if not model_loaded:
+        return "<!-- Model not available -->\n<h1>Model Loading Failed</h1>"
+    
     full_prompt = f"Create a complete single HTML file with embedded CSS and JavaScript for: {prompt}. Return only valid HTML code."
     
     inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
@@ -41,6 +51,9 @@ def generate_code(prompt):
 
 def improve_code(description, current_code):
     """Improve existing code"""
+    if not model_loaded:
+        return current_code
+    
     prompt = f"Improve this HTML code based on: {description}\n\nCurrent code:\n{current_code}\n\nReturn only the improved HTML code."
     
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")