Update app.py

app.py

@@ -3,11 +3,15 @@ import spaces
 import torch
 import torch.amp as amp
 
-from transformers import AutoModel, AutoTokenizer, pipeline
+from transformers import MistralForCausalLM, AutoTokenizer, pipeline
 
 repo_id = "appmana/Cosmos-1.0-Prompt-Upsampler-12B-Text2World-hf"
 
-model = AutoModel.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+model = MistralForCausalLM.from_pretrained(
+    repo_id,
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+)
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
 pipe = pipeline(
     "text-generation",