Update app.py for better compatibility

app.py

@@ -3,19 +3,26 @@ PlainEnglish-1B Gradio Demo App
 Interactive text generation interface for HuggingFace Spaces and ModelScope Studio.
 """
 
+import os
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_ID = "itriedcoding/PlainEnglish-1B"
+HF_MODEL_ID = "itriedcoding/PlainEnglish-1B"
+MS_MODEL_ID = "NeuraAI/PlainEnglish-1B"
 
+MODEL_ID = MS_MODEL_ID if os.environ.get("MODELSCOPE_API_TOKEN") else HF_MODEL_ID
+
+print(f"Loading model from: {MODEL_ID}")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float32,
+    device_map="auto",
     trust_remote_code=True,
 )
 model.eval()
+print("Model loaded successfully!")
 
 
 def generate_text(
@@ -30,6 +37,8 @@ def generate_text(
         return "Please enter a prompt."
 
     inputs = tokenizer(prompt, return_tensors="pt")
+    if torch.cuda.is_available():
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
     with torch.no_grad():
         outputs = model.generate(