benjamin-paine commited on
Commit
e439552
·
verified ·
1 Parent(s): 8cd32b4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -3,7 +3,7 @@ import spaces
3
  import torch
4
  import torch.amp as amp
5
 
6
- from transformers import MistralForCausalLM, AutoTokenizer, pipeline
7
 
8
  repo_id = "appmana/Cosmos-1.0-Prompt-Upsampler-12B-Text2World-hf"
9
 
@@ -12,7 +12,7 @@ model = MistralForCausalLM.from_pretrained(
12
  attn_implementation="flash_attention_2",
13
  torch_dtype=torch.bfloat16
14
  )
15
- tokenizer = AutoTokenizer.from_pretrained(repo_id)
16
  pipe = pipeline(
17
  "text-generation",
18
  model=model,
 
3
  import torch
4
  import torch.amp as amp
5
 
6
+ from transformers import MistralForCausalLM, LlamaTokenizer, pipeline
7
 
8
  repo_id = "appmana/Cosmos-1.0-Prompt-Upsampler-12B-Text2World-hf"
9
 
 
12
  attn_implementation="flash_attention_2",
13
  torch_dtype=torch.bfloat16
14
  )
15
+ tokenizer = LlamaTokenizer.from_pretrained(repo_id)
16
  pipe = pipeline(
17
  "text-generation",
18
  model=model,