Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -15,17 +15,19 @@ import scipy.io.wavfile as wav
|
|
15 |
# ---------------------------------------------------------------------
|
16 |
def load_llama_pipeline_zero_gpu(model_id: str, token: str):
|
17 |
try:
|
|
|
|
|
18 |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
|
19 |
model = AutoModelForCausalLM.from_pretrained(
|
20 |
model_id,
|
21 |
use_auth_token=token,
|
22 |
torch_dtype=torch.float16,
|
23 |
-
device_map="auto",
|
24 |
-
trust_remote_code=True
|
25 |
)
|
26 |
return pipeline("text-generation", model=model, tokenizer=tokenizer)
|
27 |
except Exception as e:
|
28 |
-
return
|
29 |
|
30 |
# ---------------------------------------------------------------------
|
31 |
# Generate Radio Script
|
@@ -115,4 +117,4 @@ with gr.Blocks() as demo:
|
|
115 |
# ---------------------------------------------------------------------
|
116 |
# Launch App
|
117 |
# ---------------------------------------------------------------------
|
118 |
-
demo.launch(debug=True
|
|
|
15 |
# ---------------------------------------------------------------------
|
16 |
def load_llama_pipeline_zero_gpu(model_id: str, token: str):
|
17 |
try:
|
18 |
+
if not torch.cuda.is_available():
|
19 |
+
raise RuntimeError("ZeroGPU is not properly initialized or GPU is unavailable.")
|
20 |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
|
21 |
model = AutoModelForCausalLM.from_pretrained(
|
22 |
model_id,
|
23 |
use_auth_token=token,
|
24 |
torch_dtype=torch.float16,
|
25 |
+
device_map="auto",
|
26 |
+
trust_remote_code=True
|
27 |
)
|
28 |
return pipeline("text-generation", model=model, tokenizer=tokenizer)
|
29 |
except Exception as e:
|
30 |
+
return f"Error loading model: {e}"
|
31 |
|
32 |
# ---------------------------------------------------------------------
|
33 |
# Generate Radio Script
|
|
|
117 |
# ---------------------------------------------------------------------
|
118 |
# Launch App
|
119 |
# ---------------------------------------------------------------------
|
120 |
+
demo.launch(debug=True)
|