Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,7 @@ import spaces
|
|
3 |
import torch
|
4 |
import torch.amp as amp
|
5 |
|
6 |
-
from transformers import MistralForCausalLM,
|
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 =
|
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,
|