local inference page with fixes to gpu with zerogpu + add accelerate for device mapping

app.py

@@ -48,11 +48,28 @@ except Exception as e:
 # Initialize Llama
 try:
     logger.info("Initializing Llama model...")
+    llm_model_id = "chuanli11/Llama-3.2-3B-Instruct-uncensored"
+    
+    # Initialize tokenizer first
+    tokenizer = AutoTokenizer.from_pretrained(llm_model_id)
+    tokenizer.use_default_system_prompt = False
+    
+    # Initialize model with proper device mapping
     if torch.cuda.is_available():
-        llm_model_id = "chuanli11/Llama-3.2-3B-Instruct-uncensored"
-        llm = AutoModelForCausalLM.from_pretrained(llm_model_id, torch_dtype=torch.float16, device_map="auto")
-        tokenizer = AutoTokenizer.from_pretrained(llm_model_id)
-        tokenizer.use_default_system_prompt = False
+        logger.info("Loading Llama model on GPU...")
+        llm = AutoModelForCausalLM.from_pretrained(
+            llm_model_id,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            load_in_8bit=True  # Use 8-bit quantization to reduce memory usage
+        )
+    else:
+        logger.info("Loading Llama model on CPU...")
+        llm = AutoModelForCausalLM.from_pretrained(
+            llm_model_id,
+            device_map={"": "cpu"},
+            low_cpu_mem_usage=True
+        )
 except Exception as e:
     logger.error(f"Error initializing Llama model: {str(e)}")
     llm = None

requirements.txt

@@ -6,4 +6,5 @@ Pillow>=10.0.0
 huggingface-hub>=0.20.0
 torch>=2.0.0
 transformers>=4.36.0
-spaces>=0.19.3 
+spaces>=0.19.3
+accelerate>=0.27.0