Update app.py

app.py

@@ -1,3 +1,4 @@
+import re
 import gradio as gr
 import spaces
 import torch
@@ -11,6 +12,12 @@ model = hf_hub_download(
 )
 llama = None
 
+def trim_upsampled_caption(prompt):
+    try:
+        return re.match(r".*[\"“](.*)[\"”]", prompt).group(1)
+    except:
+        return prompt
+
 @spaces.GPU
 def upsample(prompt):
     global llama
@@ -26,7 +33,8 @@ def upsample(prompt):
         max_tokens=512,
     )
     response = completion["choices"][0]["message"]["content"]
-    return response
+    trimmed = trim_upsampled_caption(response)
+    return response, trimmed
 
 demo = gr.Interface(
     title="NVIDIA Cosmos 🌌 Prompt Upsampler",
@@ -38,9 +46,17 @@ demo = gr.Interface(
         label="Prompt",
         value="A dog playing with a ball."
     ),
-    outputs=gr.Text(
-        label="Upsampled Prompt",
-        interactive=False
-    )
+    outputs=[
+        gr.Text(
+            label="Trimmed Upsampled Prompt",
+            interactive=False,
+            value="In a sun-drenched backyard, a playful golden retriever bounds joyfully across the lush green grass, its tail wagging with excitement. The dog, adorned with a vibrant blue collar, is captivated by a bright yellow ball, which it cradles gently in its mouth, its eyes sparkling with anticipation. The camera, held steady, captures the dog's animated expressions as it trots towards the viewer, its movements fluid and graceful. The scene is bathed in warm, golden-hour light, enhancing the vibrant colors of the dog's fur and the lush surroundings. As the dog playfully tosses the ball into the air, it leaps with enthusiasm, showcasing its agility and playful spirit. The serene backdrop of a white fence and a distant tree line frames this heartwarming moment, inviting viewers to share in the simple joy of a dog's playful antics."
+        ),
+        gr.Text(
+            label="Raw Upsampled Prompt",
+            interactive=False,
+            value="“In a sun-drenched backyard, a playful golden retriever bounds joyfully across the lush green grass, its tail wagging with excitement. The dog, adorned with a vibrant blue collar, is captivated by a bright yellow ball, which it cradles gently in its mouth, its eyes sparkling with anticipation. The camera, held steady, captures the dog's animated expressions as it trots towards the viewer, its movements fluid and graceful. The scene is bathed in warm, golden-hour light, enhancing the vibrant colors of the dog's fur and the lush surroundings. As the dog playfully tosses the ball into the air, it leaps with enthusiasm, showcasing its agility and playful spirit. The serene backdrop of a white fence and a distant tree line frames this heartwarming moment, inviting viewers to share in the simple joy of a dog's playful antics.”"
+        )
+    ]
 )
 demo.launch()