edit app.py

app.py

@@ -1,23 +1,24 @@
 import gradio as grad
 import torch
+import spaces
 from unsloth import FastLanguageModel
 from transformers import AutoTokenizer
 
 # Load your fine-tuned Phi-3 model from Hugging Face
 MODEL_NAME = "jcrissa/phi3-new-t2i"
 
-# Ensure that the device is correctly set (ZeroGPU or CPU)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+device = "cuda"
+
+@spaces.GPU
 def load_phi3_model():
     # Load the Phi-3 model and tokenizer from Hugging Face
     model, tokenizer = FastLanguageModel.from_pretrained(
         MODEL_NAME,
         max_seq_length=4096,  # Ensure it matches your fine-tuning
-        dtype=None  # Use `torch.float16` if running on GPU
+        dtype=torch.float16  # Use `torch.float16` if running on GPU
     )
     
-    # Move the model to the appropriate device (ZeroGPU or CPU)
     model.to(device)
     
     # Configure tokenizer settings