Upload app.py with huggingface_hub

app.py

@@ -5,7 +5,7 @@ Model: amkyawdev/mm-llm-tiny
 
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConfig
 import gradio as gr
 
 # Model name
@@ -14,12 +14,29 @@ MODEL_NAME = "amkyawdev/mm-llm-tiny"
 # Load model and tokenizer
 print(f"Loading model: {MODEL_NAME}...")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    # Try with device_map first
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        low_cpu_mem_usage=True
+    )
+    
+    print("Model loaded on GPU!")
+    
+except Exception as e:
+    print(f"GPU failed: {e}, trying CPU...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
+    )
+    model = model.to("cpu")
+    print("Model loaded on CPU!")
 
 # Create pipeline
 pipe = pipeline(