Update app.py

app.py

@@ -1,176 +1,51 @@
-import os
-import os.path as osp
-
 import gradio as gr
-import spaces
 import torch
-from threading import Thread
-from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer
-
-
-HEADER = ("""
-<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
-  <a href="" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
-  </a>
-  <div>
-    <h1>VideoGPT: Frontier Multimodal Foundation Models for Video Understanding</h1>
-    <h5 style="margin: 0;"></h5>
-  </div>
-</div>
-""")
-
-device = "cuda"
-model = AutoModelForCausalLM.from_pretrained(
-    "OpenGVLab/InternVideo2_5_Chat_8B",
-    trust_remote_code=True,
-    torch_dtype=torch.bfloat16,
-    attn_implementation="flash_attention_2",
-)
-model.to(device)
-processor = AutoProcessor.from_pretrained("DAMO-NLP-SG/VideoLLaMA3-7B", trust_remote_code=True)
-
-
-example_dir = "./examples"
-image_formats = ("png", "jpg", "jpeg")
-video_formats = ("mp4",)
-
-image_examples, video_examples = [], []
-if example_dir is not None:
-    example_files = [
-        osp.join(example_dir, f) for f in os.listdir(example_dir)
-    ]
-    for example_file in example_files:
-        if example_file.endswith(image_formats):
-            image_examples.append([example_file])
-        elif example_file.endswith(video_formats):
-            video_examples.append([example_file])
-
-
-def _on_video_upload(messages, video):
-        if video is not None:
-            # messages.append({"role": "user", "content": gr.Video(video)})
-            messages.append({"role": "user", "content": {"path": video}})
-        return messages, None
-    
-def _on_image_upload(messages, image):
-    if image is not None:
-        # messages.append({"role": "user", "content": gr.Image(image)})
-        messages.append({"role": "user", "content": {"path": image}})
-    return messages, None
-
-def _on_text_submit(messages, text):
-    messages.append({"role": "user", "content": text})
-    return messages, ""
-
-@spaces.GPU(duration=120)
-def _predict(messages, input_text, do_sample, temperature, top_p, max_new_tokens,
-             fps, max_frames):
-    if len(input_text) > 0:
-        messages.append({"role": "user", "content": input_text})
-    new_messages = []
-    contents = []
-    for message in messages:
-        if message["role"] == "assistant":
-            if len(contents):
-                new_messages.append({"role": "user", "content": contents})
-                contents = []
-            new_messages.append(message)
-        elif message["role"] == "user":
-            if isinstance(message["content"], str):
-                contents.append(message["content"])
-            else:
-                media_path = message["content"][0]
-                if media_path.endswith(video_formats):
-                    contents.append({"type": "video", "video": {"video_path": media_path, "fps": fps, "max_frames": max_frames}})
-                elif media_path.endswith(image_formats):
-                    contents.append({"type": "image", "image": {"image_path": media_path}})
-                else:
-                    raise ValueError(f"Unsupported media type: {media_path}")
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
 
-    if len(contents):
-        new_messages.append({"role": "user", "content": contents})
+# Model name
+model_name = "OpenGVLab/InternVideo2_5_Chat_8B"
 
-    if len(new_messages) == 0 or new_messages[-1]["role"] != "user":
-        return messages
-
-    generation_config = {
-        "do_sample": do_sample,
-        "temperature": temperature,
-        "top_p": top_p,
-        "max_new_tokens": max_new_tokens
-    }
-
-    inputs = processor(
-        conversation=new_messages,
-        add_system_prompt=True,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    )
-    inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
-    if "pixel_values" in inputs:
-        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
-
-    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {
-        **inputs,
-        **generation_config,
-        "streamer": streamer,
-    }
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    messages.append({"role": "assistant", "content": ""})
-    for token in streamer:
-        messages[-1]['content'] += token
-        yield messages
-
-
-with gr.Blocks() as interface:
-    gr.HTML(HEADER)
-    with gr.Row():
-        chatbot = gr.Chatbot(type="messages", elem_id="chatbot", height=835)
-
-        with gr.Column():
-            with gr.Tab(label="Input"):
-
-                with gr.Row():
-                    input_video = gr.Video(sources=["upload"], label="Upload Video")
-                    input_image = gr.Image(sources=["upload"], type="filepath", label="Upload Image")
-                
-                input_text = gr.Textbox(label="Input Text", placeholder="Type your message here and press enter to submit")
+# Set Hugging Face cache directory
+import os
+os.environ["HF_HOME"] = "/home/user/.cache/huggingface"  # Change if needed
 
-                submit_button = gr.Button("Generate")
+# Increase timeout for downloads
+os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"  # 5 minutes
 
-                gr.Examples(examples=[
-                    [f"examples/bear.mp4", "What is unusual in the video?"],
-                    [f"examples/dog.mp4", "Please describe the video in detail."],
-                    [f"examples/exercise.mp4", "What is the man doing in the video?"],
-                ], inputs=[input_video, input_text], label="Video examples")
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-            with gr.Tab(label="Configure"):
-                with gr.Accordion("Generation Config", open=True):
-                    do_sample = gr.Checkbox(value=True, label="Do Sample")
-                    temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Temperature")
-                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top P")
-                    max_new_tokens = gr.Slider(minimum=0, maximum=4096, value=2048, step=1, label="Max New Tokens")
+# Detect device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-                with gr.Accordion("Video Config", open=True):
-                    fps = gr.Slider(minimum=0.0, maximum=10.0, value=1, label="FPS")
-                    max_frames = gr.Slider(minimum=0, maximum=256, value=180, step=1, label="Max Frames")
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None
+)
 
-    input_video.change(_on_video_upload, [chatbot, input_video], [chatbot, input_video])
-    input_image.change(_on_image_upload, [chatbot, input_image], [chatbot, input_image])
-    input_text.submit(_on_text_submit, [chatbot, input_text], [chatbot, input_text])
-    submit_button.click(
-        _predict,
-        [
-            chatbot, input_text, do_sample, temperature, top_p, max_new_tokens,
-            fps, max_frames
-        ],
-        [chatbot],
-    )
+# Move model to device
+model.to(device)
 
+# Define inference function
+def chat_with_model(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    output = model.generate(**inputs, max_length=200)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+# Create Gradio UI
+demo = gr.Interface(
+    fn=chat_with_model,
+    inputs=gr.Textbox(placeholder="Type your prompt here..."),
+    outputs="text",
+    title="InternVideo2.5 Chatbot",
+    description="A chatbot powered by InternVideo2_5_Chat_8B.",
+    theme="compact"
+)
 
+# Run the Gradio app
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()