Update app.py

app.py

@@ -6,12 +6,11 @@ from decord import cpu, VideoReader, bridge
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import BitsAndBytesConfig
 
-MODEL_PATH = "THUDM/cogvlm2-llama3-caption"
+MODEL_PATH = "THUDM/cogvlm2-video-llama3-chat"
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16
 
 
-
 def get_step_info(step_number):
     """Returns detailed information about a manufacturing step."""
     step_details = {
@@ -126,89 +125,6 @@ def get_step_info(step_number):
     
     return step_details.get(step_number, {"Error": "Invalid step number. Please provide a valid step number."})
 
-
-
-def load_video(video_data, strategy='chat'):
-    """Loads and processes video data into a format suitable for model input."""
-    bridge.set_bridge('torch')
-    num_frames = 24
-    
-    if isinstance(video_data, str):
-        decord_vr = VideoReader(video_data, ctx=cpu(0))
-    else:  
-        decord_vr = VideoReader(io.BytesIO(video_data), ctx=cpu(0))
-    
-    frame_id_list = []
-    total_frames = len(decord_vr)
-    timestamps = [i[0] for i in decord_vr.get_frame_timestamp(np.arange(total_frames))]
-    max_second = round(max(timestamps)) + 1
-    
-    for second in range(max_second):
-        closest_num = min(timestamps, key=lambda x: abs(x - second))
-        index = timestamps.index(closest_num)
-        frame_id_list.append(index)
-        if len(frame_id_list) >= num_frames:
-            break
-
-    video_data = decord_vr.get_batch(frame_id_list)
-    video_data = video_data.permute(3, 0, 1, 2)
-    return video_data
-
-def load_model():
-    """Loads the pre-trained model and tokenizer with quantization configurations."""
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=TORCH_TYPE,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4"
-    )
-    
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_PATH,
-        torch_dtype=TORCH_TYPE,
-        trust_remote_code=True,
-        quantization_config=quantization_config,
-        device_map="auto"
-    ).eval()
-    
-    return model, tokenizer
-
-def predict(prompt, video_data, temperature, model, tokenizer):
-    """Generates predictions based on the video and textual prompt."""
-    video = load_video(video_data, strategy='chat')
-    
-    inputs = model.build_conversation_input_ids(
-        tokenizer=tokenizer,
-        query=prompt,
-        images=[video],
-        history=[],
-        template_version='chat'
-    )
-    
-    inputs = {
-        'input_ids': inputs['input_ids'].unsqueeze(0).to(DEVICE),
-        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(DEVICE),
-        'attention_mask': inputs['attention_mask'].unsqueeze(0).to(DEVICE),
-        'images': [[inputs['images'][0].to(DEVICE).to(TORCH_TYPE)]],
-    }
-    
-    gen_kwargs = {
-        "max_new_tokens": 2048,
-        "pad_token_id": 128002,
-        "top_k": 1,
-        "do_sample": False,
-        "top_p": 0.1,
-        "temperature": temperature,
-    }
-    
-    with torch.no_grad():
-        outputs = model.generate(**inputs, **gen_kwargs)
-        outputs = outputs[:, inputs['input_ids'].shape[1]:]
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return response
-
 def get_analysis_prompt(step_number):
     """Constructs the prompt for analyzing delay reasons based on the selected step."""
     step_info = get_step_info(step_number)
@@ -254,10 +170,43 @@ Output:
     No person available to collect tire
 """
 
-
-
 model, tokenizer = load_model()
 
+def predict(prompt, video_data, temperature, model, tokenizer):
+    """Generates predictions based on the video and textual prompt."""
+    video = load_video(video_data, strategy='chat')
+    
+    inputs = model.build_conversation_input_ids(
+        tokenizer=tokenizer,
+        query=prompt,
+        images=[video],
+        history=[],
+        template_version='chat'
+    )
+    
+    inputs = {
+        'input_ids': inputs['input_ids'].unsqueeze(0).to(DEVICE),
+        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(DEVICE),
+        'attention_mask': inputs['attention_mask'].unsqueeze(0).to(DEVICE),
+        'images': [[inputs['images'][0].to(DEVICE).to(TORCH_TYPE)]],
+    }
+    
+    gen_kwargs = {
+        "max_new_tokens": 2048,
+        "pad_token_id": 128002,
+        "top_k": 1,
+        "do_sample": False,
+        "top_p": 0.1,
+        "temperature": temperature,
+    }
+    
+    with torch.no_grad():
+        outputs = model.generate(**inputs, **gen_kwargs)
+        outputs = outputs[:, inputs['input_ids'].shape[1]:]
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return response
+
 def inference(video, step_number):
     """Analyzes video to predict possible issues based on the manufacturing step."""
     try:
@@ -315,4 +264,4 @@ def create_interface():
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.queue().launch(share=True)
+    demo.queue().launch(share=True)