Bils commited on
Commit
7bbdf94
·
verified ·
1 Parent(s): f7e7de6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -4
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", # Use device map to offload computations
24
- trust_remote_code=True # Enables execution of remote code for Zero GPU
25
  )
26
  return pipeline("text-generation", model=model, tokenizer=tokenizer)
27
  except Exception as e:
28
- return str(e)
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, share=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)