Update app.py

app.py

@@ -79,125 +79,132 @@ def progress_bar_html(label: str) -> str:
     '''
 
 @spaces.GPU
-def generate(text: str, files: list,
-             max_new_tokens: int = 1024,
-             temperature: float = 0.6,
-             top_p: float = 0.9,
-             top_k: int = 50,
-             repetition_penalty: float = 1.2):
+def generate_image(text: str, image: Image.Image,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
     """
-    Generates responses using the Qwen2VL model for image and video inputs.
-    - If images are provided, performs image inference.
-    - If videos are provided, performs video inference by downsampling to frames.
+    Generates responses using the Cosmos-Reason1 model for image input.
     """
-    if not files:
-        yield "Please upload an image or video for inference."
+    if image is None:
+        yield "Please upload an image."
         return
 
-    # Determine if the files are images or videos
-    image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
-    video_files = [f for f in files if f.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))]
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": text},
+        ]
+    }]
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(
+        text=[prompt_full],
+        images=[image],
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=MAX_INPUT_TOKEN_LENGTH
+    ).to("cuda")
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
+    thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    yield progress_bar_html("Processing image with Cosmos-Reason1")
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "")
+        time.sleep(0.01)
+        yield buffer
 
-    if image_files and video_files:
-        yield "Please upload either images or videos, not both."
+@spaces.GPU
+def generate_video(text: str, video_path: str,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
+    """
+    Generates responses using the Cosmos-Reason1 model for video input.
+    """
+    if video_path is None:
+        yield "Please upload a video."
         return
 
-    if image_files:
-        # Image inference
-        images = [load_image(image) for image in image_files]
-        messages = [{
-            "role": "user",
-            "content": [
-                *[{"type": "image", "image": image} for image in images],
-                {"type": "text", "text": text},
-            ]
-        }]
-        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(
-            text=[prompt_full],
-            images=images,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=MAX_INPUT_TOKEN_LENGTH
-        ).to("cuda")
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
-        thread.start()
-        buffer = ""
-        yield progress_bar_html("Processing images with cosmos-reasoning")
-        for new_text in streamer:
-            buffer += new_text
-            buffer = buffer.replace("<|im_end|>", "")
-            time.sleep(0.01)
-            yield buffer
-    elif video_files:
-        # Video inference
-        video_path = video_files[0]  # Assuming only one video is uploaded
-        frames = downsample_video(video_path)
-        messages = [
-            {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-            {"role": "user", "content": [{"type": "text", "text": text}]}
-        ]
-        # Append each frame with its timestamp.
-        for frame in frames:
-            image, timestamp = frame
-            image_path = f"video_frame_{uuid.uuid4().hex}.png"
-            image.save(image_path)
-            messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
-            messages[1]["content"].append({"type": "image", "url": image_path})
-        inputs = processor.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_dict=True,
-            return_tensors="pt",
-            truncation=True,
-            max_length=MAX_INPUT_TOKEN_LENGTH
-        ).to("cuda")
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {
-            **inputs,
-            "streamer": streamer,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-        }
-        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
-        thread.start()
-        buffer = ""
-        yield progress_bar_html("Processing video with cosmos-reasoning")
-        for new_text in streamer:
-            buffer += new_text
-            buffer = buffer.replace("<|im_end|>", "")
-            time.sleep(0.01)
-            yield buffer
-    else:
-        yield "Unsupported file type. Please upload images or videos."
+    frames = downsample_video(video_path)
+    messages = [
+        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
+        {"role": "user", "content": [{"type": "text", "text": text}]}
+    ]
+    # Append each frame with its timestamp.
+    for frame in frames:
+        image, timestamp = frame
+        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
+        messages[1]["content"].append({"type": "image", "image": image})
+    inputs = processor.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt",
+        truncation=True,
+        max_length=MAX_INPUT_TOKEN_LENGTH
+    ).to("cuda")
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {
+        **inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "repetition_penalty": repetition_penalty,
+    }
+    thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    yield progress_bar_html("Processing video with Cosmos-Reason1")
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "")
+        time.sleep(0.01)
+        yield buffer
 
 # Create the Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# **cosmos-reason1 by nvidia**")
+    gr.Markdown("# **Cosmos-Reason1 by NVIDIA**")
     with gr.Row():
         with gr.Column():
-            text_input = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
-            file_input = gr.File(label="Upload Image or Video", file_types=["image", "video"], file_count="multiple")
-            max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
-            temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
-            top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
-            top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
-            repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
-            submit_btn = gr.Button("Submit")
+            with gr.Tabs():
+                with gr.TabItem("Image Inference"):
+                    image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    image_upload = gr.Image(type="pil", label="Upload Image")
+                    image_submit = gr.Button("Submit")
+                with gr.TabItem("Video Inference"):
+                    video_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    video_upload = gr.Video(label="Upload Video")
+                    video_submit = gr.Button("Submit")
+            with gr.Accordion("Advanced options", open=False):
+                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
+                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
         with gr.Column():
             output = gr.Textbox(label="Output", interactive=False)
 
-    submit_btn.click(
-        fn=generate,
-        inputs=[text_input, file_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+    image_submit.click(
+        fn=generate_image,
+        inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
+    )
+    video_submit.click(
+        fn=generate_video,
+        inputs=[video_query, video_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
         outputs=output
     )