Update app.py

app.py

@@ -3,25 +3,37 @@ import spaces
 import torch
 import torch.amp as amp
 
-from transformers import pipeline
+from transformers import AutoModel, AutoTokenizer, pipeline
 
-upsampler = pipeline(
-    "text-generation", 
-    "appmana/Cosmos-1.0-Prompt-Upsampler-12B-Text2World-hf",
+repo_id = "appmana/Cosmos-1.0-Prompt-Upsampler-12B-Text2World-hf"
+
+model = AutoModel.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
     torch_dtype=torch.bfloat16
 )
 
 @spaces.GPU
 def upsample(prompt):
-    return upsampler(
+    template = tokenizer.apply_chat_template(
         [{"role": "user", "content": f"Upsample the short caption to a long caption: {prompt}"}],
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    response = pipe(
+        template,
         min_p=0.01,
         top_p=0.95,
         top_k=40,
+        do_sample=True,
         temperature=0.2,
         max_new_tokens=512,
-        do_sample=True
-    )[0]["generated_text"][-1]["content"]
+        pad_token_id=tokenizer.eos_token_id
+    )
+    return response[0]["generated_text"]
 
 demo = gr.Interface(
     title="NVIDIA Cosmos 🌌 Prompt Upsampler",