fix

process.py

@@ -1,32 +1,28 @@
 import io
-import spaces
-
 import argparse
 import numpy as np
 import torch
 from decord import cpu, VideoReader, bridge
 from transformers import AutoModelForCausalLM, AutoTokenizer
-
+from transformers import BitsAndBytesConfig
 
 MODEL_PATH = "THUDM/cogvlm2-llama3-caption"
-
 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
 
 parser = argparse.ArgumentParser(description="CogVLM2-Video CLI Demo")
-parser.add_argument('--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=0)
+parser.add_argument('--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=4)
 args = parser.parse_args([])
 
-
 def load_video(video_data, strategy='chat'):
     bridge.set_bridge('torch')
     mp4_stream = video_data
     num_frames = 24
     decord_vr = VideoReader(io.BytesIO(mp4_stream), ctx=cpu(0))
-
     frame_id_list = None
     total_frames = len(decord_vr)
+    
     if strategy == 'base':
         clip_end_sec = 60
         clip_start_sec = 0
@@ -45,11 +41,18 @@ def load_video(video_data, strategy='chat'):
             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
 
+# Configure quantization
+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,
@@ -59,11 +62,14 @@ tokenizer = AutoTokenizer.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
     torch_dtype=TORCH_TYPE,
-    trust_remote_code=True
-).eval().to(DEVICE)
+    trust_remote_code=True,
+    quantization_config=quantization_config,
+    device_map="auto"
+).eval()
 
-@spaces.GPU
-def predict(prompt, video, temperature, strategy):
+def predict(prompt, video_data, temperature):
+    strategy = 'chat'
+    video = load_video(video_data, strategy=strategy)
     history = []
     query = prompt
     inputs = model.build_conversation_input_ids(
@@ -73,31 +79,34 @@ def predict(prompt, video, temperature, strategy):
         history=history,
         template_version=strategy
     )
+    
     inputs = {
-        'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'),
-        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to('cuda'),
-        'attention_mask': inputs['attention_mask'].unsqueeze(0).to('cuda'),
-        'images': [[inputs['images'][0].to('cuda').to(TORCH_TYPE)]],
+        '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": True,
+        "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, prompt):
+def test(video, prompt):
     temperature = 0.1
-    video = open(video, 'rb').read()
-    strategy = 'base'
-    video_data = load_video(video, strategy=strategy)
-    response = predict(prompt, video_data, temperature, strategy)
-    return response
+    video_data = open(video, 'rb').read()
+    response = predict(prompt, video_data, temperature)
+    print(response)
+
+if __name__ == '__main__':
+    test()

requirements.txt

@@ -1,8 +1,25 @@
-argparse
-numpy
-decord
-spaces
-transformers==4.44.2
-einops==0.8.0 
-torchvision==0.16.1
-pytorchvideo
+decord>=0.6.0
+#根据https://download.pytorch.org/whl/torch/，python版本为[3.8,3.11]
+torch==2.1.0
+torchvision== 0.16.0
+pytorchvideo==0.1.5
+xformers
+transformers==4.42.4
+#git+https://github.com/huggingface/transformers.git
+huggingface-hub>=0.23.0
+pillow
+chainlit>=1.0
+pydantic>=2.7.1
+timm>=0.9.16
+openai>=1.30.1
+loguru>=0.7.2
+pydantic>=2.7.1
+einops
+sse-starlette>=2.1.0
+flask
+gunicorn
+gevent
+requests
+gradio
+accelerate
+bitsandbytes>=0.39.0