edit app.py

app.py

@@ -6,13 +6,21 @@ 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")
+
 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
     )
     
+    # Move the model to the appropriate device (ZeroGPU or CPU)
+    model.to(device)
+    
+    # Configure tokenizer settings
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "left"
     
@@ -22,9 +30,11 @@ phi3_model, phi3_tokenizer = load_phi3_model()
 
 # Function to generate text using Phi-3
 def generate(plain_text):
-    input_ids = phi3_tokenizer(plain_text.strip(), return_tensors="pt").input_ids
+    # Tokenize input text and move to the device
+    input_ids = phi3_tokenizer(plain_text.strip(), return_tensors="pt").input_ids.to(device)
     eos_id = phi3_tokenizer.eos_token_id
     
+    # Generate the output from the model
     outputs = phi3_model.generate(
         input_ids, 
         do_sample=True,  
@@ -36,6 +46,7 @@ def generate(plain_text):
         length_penalty=-1.0
     )
     
+    # Decode and return the generated text
     output_text = phi3_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return output_text.strip()
 
@@ -43,6 +54,7 @@ def generate(plain_text):
 txt = grad.Textbox(lines=1, label="Input Text", placeholder="Enter your prompt")
 out = grad.Textbox(lines=1, label="Generated Text")
 
+# Launch Gradio Interface with ZeroGPU-compatible setup
 grad.Interface(
     fn=generate,
     inputs=txt,