edit app.py

app.py

@@ -7,56 +7,67 @@ from transformers import AutoTokenizer
 # Load your fine-tuned Phi-3 model from Hugging Face
 MODEL_NAME = "jcrissa/phi3-new-t2i"
 
+# Check if CUDA is available, otherwise fall back to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-device = "cuda"
-
-@spaces.GPU
+# Function to load the Phi-3 model and tokenizer
+@spaces.GPU  # Reintroduced spaces.GPU decorator for GPU setup
 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=torch.float16  # Use `torch.float16` if running on GPU
-    )
-    
-    model.to(device)
-    
-    # Configure tokenizer settings
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.padding_side = "left"
-    
-    return model, tokenizer
+    try:
+        # 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=torch.float16 if device == "cuda" else torch.float32  # Use `float16` for GPU, `float32` for CPU
+        )
+        model.to(device)
+
+        # Configure tokenizer settings
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.padding_side = "left"
+        
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None
 
+# Load the model and tokenizer, ensure error handling
 phi3_model, phi3_tokenizer = load_phi3_model()
 
+if phi3_model is None or phi3_tokenizer is None:
+    raise RuntimeError("Model and tokenizer could not be loaded. Please check the Hugging Face model path or network connection.")
+
 # Function to generate text using Phi-3
 def generate(plain_text):
-    # 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,  
-        max_new_tokens=75, 
-        num_beams=8, 
-        num_return_sequences=1, 
-        eos_token_id=eos_id, 
-        pad_token_id=eos_id,
-        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()
+    try:
+        # 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,  
+            max_new_tokens=75, 
+            num_beams=8, 
+            num_return_sequences=1, 
+            eos_token_id=eos_id, 
+            pad_token_id=eos_id,
+            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()
+    except Exception as e:
+        return f"Error during text generation: {e}"
 
 # Setup Gradio Interface
 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(
+gr.Interface(
     fn=generate,
     inputs=txt,
     outputs=out,