cpu support

app.py

@@ -22,7 +22,7 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = AutoModelForCausalLM.from_pretrained(
   model_id, 
-  device_map="cuda", 
+  device_map=DEVICE, 
   trust_remote_code=True, 
   torch_dtype="auto", 
 #  _attn_implementation='flash_attention_2',
@@ -43,7 +43,7 @@ def phi_ocr(image_url):
     }]
 
     prompt = processor.tokenizer.apply_chat_template(prompt_message, tokenize=False, add_generation_prompt=True)
-    inputs = processor(prompt, [image], return_tensors="pt").to("cuda") 
+    inputs = processor(prompt, [image], return_tensors="pt").to(DEVICE) 
 
     generation_args = { 
         "max_new_tokens": 1024,