the latest version based on inference.py

handler.py

@@ -1,142 +1,77 @@
 import os
-import torch
 import base64
-from PIL import Image
-from io import BytesIO
-from typing import Dict, Any
-from transformers import LlamaTokenizer, GenerationConfig
-from robohusky.model.modeling_husky_embody2 import HuskyForConditionalGeneration
-from decord import VideoReader, cpu
-import torchvision.transforms as T
-from torchvision.transforms.functional import InterpolationMode
 import tempfile
+from inference import Chat, get_conv_template
+import torch
 
-DEFAULT_IMG_START_TOKEN = "<img>"
-DEFAULT_IMG_END_TOKEN = "</img>"
-DEFAULT_VIDEO_START_TOKEN = "<vid>"
-DEFAULT_VIDEO_END_TOKEN = "</vid>"
+def save_base64_to_tempfile(base64_str, suffix):
+    header_removed = base64_str
+    # 去除可能的data:image/...;base64,前缀
+    if ',' in base64_str:
+        header_removed = base64_str.split(',', 1)[1]
 
-class EndpointHandler:
-    def __init__(self, model_path: str = "."):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.tokenizer = LlamaTokenizer.from_pretrained(model_path, use_fast=False)
-        self.model = HuskyForConditionalGeneration.from_pretrained(
-            model_path, torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
-        ).to(self.device).eval()
-
-        self.gen_config = GenerationConfig(
-            bos_token_id=1,
-            do_sample=False,
-            # temperature=0.7,
-            max_new_tokens=10240
+    data = base64.b64decode(header_removed)
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+    tmp.write(data)
+    tmp.close()
+    return tmp.name
 
+class EndpointHandler:
+    def __init__(self, model_path: str):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.chat = Chat(
+            model_path=model_path,
+            device=device,
+            num_gpus=1,
+            max_new_tokens=1024,
+            load_8bit=False,
         )
+        self.vision_feature = None
+        self.modal_type = "text"
+        self.chat.conv = get_conv_template("husky").copy()
+
+    def __call__(self, data: dict) -> dict:
+        # reset conversation if specified
+        if data.get("clear_history"):
+            self.chat.conv = get_conv_template("husky").copy()
+            self.vision_feature = None
+            self.modal_type = "text"
+
+        prompt = data.get("inputs", "")
+        image_input = data.get("image", None)
+        video_input = data.get("video", None)
+
+        # 判断image输入是路径还是base64字符串
+        if image_input:
+            if os.path.exists(image_input):
+                # 直接用路径
+                self.vision_feature = self.chat.get_image_embedding(image_input)
+            else:
+                # base64字符串，保存临时文件再处理
+                tmp_path = save_base64_to_tempfile(image_input, suffix=".jpg")
+                self.vision_feature = self.chat.get_image_embedding(tmp_path)
+                os.unlink(tmp_path)  # 删除临时文件
+            self.modal_type = "image"
+            self.chat.conv = get_conv_template("husky").copy()
+
+        elif video_input:
+            if os.path.exists(video_input):
+                self.vision_feature = self.chat.get_video_embedding(video_input)
+            else:
+                tmp_path = save_base64_to_tempfile(video_input, suffix=".mp4")
+                self.vision_feature = self.chat.get_video_embedding(tmp_path)
+                os.unlink(tmp_path)
+            self.modal_type = "video"
+            self.chat.conv = get_conv_template("husky").copy()
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        inputs = self.preprocess(data)
-        prediction = self.inference(inputs)
-        return self.postprocess(prediction)
-
-    def preprocess(self, request: Dict[str, Any]) -> Dict[str, Any]:
-        prompt = request["inputs"]
-        image_b64 = request.get("image", None)
-        video_b64 = request.get("video", None)
-
-        pixel_values = None
-
-        if image_b64:
-            image_bytes = base64.b64decode(image_b64)
-            pixel_values = self._load_image(image_bytes).unsqueeze(0)  # [1, 3, 224, 224]
-            if self.device == "cuda":
-                pixel_values = pixel_values.half()
-            pixel_values = pixel_values.to(self.device)
-            prompt = prompt.replace("<image>", DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN)
-
-        elif video_b64:
-            video_bytes = base64.b64decode(video_b64)
-            pixel_values = self._load_video(video_bytes)
-            if self.device == "cuda":
-                pixel_values = pixel_values.half()
-            pixel_values = pixel_values.to(self.device)
-            prompt = prompt.replace("<video>", DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN)
-
-        return {
-            "prompt": prompt,
-            "pixel_values": pixel_values
-        }
-
-    def inference(self, inputs: Dict[str, Any]) -> str:
-        prompt = inputs["prompt"]
-        pixel_values = inputs["pixel_values"]
-
-        model_inputs = self.tokenizer([prompt], return_tensors="pt")
-        model_inputs.pop("token_type_ids", None)
-        model_inputs = {k: v.to(self.device) for k, v in model_inputs.items()}
-
-        if pixel_values is not None:
-            output = self.model.generate(
-                **model_inputs,
-                pixel_values=pixel_values,
-                max_new_tokens=self.gen_config.max_new_tokens,  # 👈 显式传入
-                generation_config=self.gen_config,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
         else:
-            output = self.model.language_model.generate(
-                **model_inputs,
-                generation_config=self.gen_config,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
-        # 🧠 打印 debug 信息
-        generated_ids = output.sequences[0]
-        print("📍生成的 token ids：", generated_ids.tolist())
-        raw_text = self.tokenizer.decode(generated_ids, skip_special_tokens=False)
-        clean_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
-        print("🧾 带特殊符号的输出：", raw_text)
-        print("✅ 去掉特殊符号的输出：", clean_text)
+            self.modal_type = "text"
+            self.vision_feature = None
 
-        return clean_text  # 返回干净版本
-    def postprocess(self, output: str) -> Dict[str, str]:
-        return {"output": output.strip()}
+        conversations = self.chat.ask(prompt, self.chat.conv, modal_type=self.modal_type)
+        output = self.chat.answer(conversations, self.vision_feature, modal_type=self.modal_type)
 
-    def _load_image(self, image_bytes: bytes) -> torch.Tensor:
-        image = Image.open(BytesIO(image_bytes)).convert('RGB')
-        crop_pct = 224 / 256
-        size = int(224 / crop_pct)
-        transform = T.Compose([
-            T.Resize(size, interpolation=InterpolationMode.BICUBIC),
-            T.CenterCrop(224),
-            T.ToTensor(),
-            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
-        ])
-        return transform(image)
+        # 更新对话历史
+        self.chat.conv.messages[-1][1] = output.strip()
 
-    def _load_video(self, video_bytes: bytes, num_segments=8) -> torch.Tensor:
-        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
-            tmpfile.write(video_bytes)
-            video_path = tmpfile.name
-
-        vr = VideoReader(video_path, ctx=cpu(0))
-        total_frames = len(vr)
-        indices = self.get_index(total_frames, num_segments)
-        frames = [Image.fromarray(vr[i].asnumpy()) for i in indices]
-
-        transform = T.Compose([
-            T.Resize(224, interpolation=InterpolationMode.BICUBIC),
-            T.CenterCrop(224),
-            T.ToTensor(),
-            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
-        ])
-        processed = [transform(frame) for frame in frames]  # each: [3, 224, 224]
-        video_tensor = torch.stack(processed, dim=0)  # [T, 3, 224, 224]
-        video_tensor = video_tensor.permute(1, 0, 2, 3)  # [3, T, 224, 224]
-        video_tensor = video_tensor.unsqueeze(0)  # [1, 3, T, 224, 224] ✅
-        return video_tensor
-
-    def get_index(self, num_frames: int, num_segments: int):
-        if num_frames < num_segments:
-            return list(range(num_frames)) + [num_frames - 1] * (num_segments - num_frames)
-        interval = num_frames / num_segments
-        return [int(interval * i) for i in range(num_segments)]
+        return {"output": output.strip()}

inference.py

@@ -0,0 +1,415 @@
+"""
+srun -p INTERN2 --job-name='husky_multi_test' --gres=gpu:1 --cpus-per-task=8 --quotatype="auto" python -u demo/inference_new.py
+"""
+
+import abc
+from typing import Optional
+
+import os
+import requests
+from PIL import Image
+from io import BytesIO
+
+import torch
+import torchvision.transforms as T
+from peft import PeftModel
+from torchvision.transforms.functional import InterpolationMode
+
+from transformers import (
+    LlamaTokenizer,
+    GenerationConfig,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
+
+from robohusky.model.modeling_husky_embody2 import HuskyForConditionalGeneration
+
+from robohusky.conversation import (
+    conv_templates,
+    get_conv_template,
+)
+
+from robohusky.video_transformers import (
+    GroupNormalize,
+    GroupScale,
+    GroupCenterCrop,
+    Stack,
+    ToTorchFormatTensor,
+    get_index,
+)
+
+from robohusky.compression import compress_module
+from decord import VideoReader, cpu
+
+# import deepspeed
+
+IGNORE_INDEX = -100
+DEFAULT_UNK_TOKEN = "<unk>"
+DEFAULT_IMG_START_TOKEN = "<img>"
+DEFAULT_IMG_END_TOKEN = "</img>"
+
+DEFAULT_VIDEO_START_TOKEN = "<vid>"
+DEFAULT_VIDEO_END_TOKEN = "</vid>"
+
+def get_gpu_memory(max_gpus=None):
+    gpu_memory = []
+    num_gpus = (
+        torch.cuda.device_count()
+        if max_gpus is None
+        else min(max_gpus, torch.cuda.device_count())
+    )
+
+    for gpu_id in range(num_gpus):
+        with torch.cuda.device(gpu_id):
+            device = torch.cuda.current_device()
+            gpu_properties = torch.cuda.get_device_properties(device)
+            total_memory = gpu_properties.total_memory / (1024 ** 3)
+            allocated_memory = torch.cuda.memory_allocated() / (1024 ** 3)
+            available_memory = total_memory - allocated_memory
+            gpu_memory.append(available_memory)
+    return gpu_memory
+
+def load_model(
+        model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, lora_weights=None
+):
+    if device == "cpu":
+        kwargs = {}
+    elif device == "cuda":
+        kwargs = {"torch_dtype": torch.float16}
+        if num_gpus == "auto":
+            kwargs["device_map"] = "auto"
+        else:
+            num_gpus = int(num_gpus)
+            if num_gpus != 1:
+                kwargs["device_map"] = "auto"
+                if max_gpu_memory is None:
+                    kwargs[
+                        "device_map"
+                    ] = "sequential"  # This is important for not the same VRAM sizes
+                    available_gpu_memory = get_gpu_memory(num_gpus)
+                    kwargs["max_memory"] = {
+                        i: str(int(available_gpu_memory[i] * 0.85)) + "GiB"
+                        for i in range(num_gpus)
+                    }
+                else:
+                    kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)}
+    else:
+        raise ValueError(f"Invalid device: {device}")
+
+    tokenizer = LlamaTokenizer.from_pretrained(
+        model_path, use_fast=False)
+
+    if lora_weights is None:
+        model = HuskyForConditionalGeneration.from_pretrained(
+            model_path, low_cpu_mem_usage=True, **kwargs
+        )
+    else:
+        kwargs["device_map"] = "auto"
+        model = HuskyForConditionalGeneration.from_pretrained(
+            model_path, low_cpu_mem_usage=True, **kwargs
+        )
+        model.language_model = PeftModel.from_pretrained(
+            model.language_model,
+            lora_weights,
+            **kwargs
+        )
+
+    if load_8bit:
+        compress_module(model, device)
+
+    if (device == "cuda" and num_gpus == 1) or device == "mps":
+        model.to(device)
+
+    model = model.eval()
+    return model, tokenizer
+
+def load_image(image_file, input_size=224):
+    if image_file.startswith('http') or image_file.startswith('https'):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        image = Image.open(image_file).convert('RGB')
+
+    crop_pct = 224 / 256
+    size = int(input_size / crop_pct)
+    transform = T.Compose([
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+        T.Resize(size, interpolation=InterpolationMode.BICUBIC),
+        T.CenterCrop(input_size),
+        T.ToTensor(),
+        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
+    ])
+    image = transform(image)
+    return image
+
+def load_video(video_path, num_segments=8):
+    vr = VideoReader(video_path, ctx=cpu(0))
+    num_frames = len(vr)
+    frame_indices = get_index(num_frames, num_segments)
+
+    # transform
+    crop_size = 224
+    scale_size = 224
+    input_mean = [0.48145466, 0.4578275, 0.40821073]
+    input_std = [0.26862954, 0.26130258, 0.27577711]
+
+    transform = T.Compose([
+        GroupScale(int(scale_size), interpolation=InterpolationMode.BICUBIC),
+        GroupCenterCrop(crop_size),
+        Stack(),
+        ToTorchFormatTensor(),
+        GroupNormalize(input_mean, input_std)
+    ])
+
+    images_group = list()
+    for frame_index in frame_indices:
+        img = Image.fromarray(vr[frame_index].asnumpy())
+        images_group.append(img)
+    video = transform(images_group)
+    return video
+
+class StoppingCriteriaSub(StoppingCriteria):
+
+    def __init__(self, stops, encounters=1):
+        super().__init__()
+        self.stops = stops
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
+        for stop in self.stops:
+            if torch.all((stop == input_ids[0][-len(stop):])).item():
+                return True
+
+        return False
+
+@torch.inference_mode()
+def generate_stream(
+        model, tokenizer, image_processor, params, device
+):
+    prompt = params["prompt"]
+    images = params.get("images", None)
+    videos = params.get("videos", None)
+    temperature = float(params.get("temperature", 0.7))
+    max_new_tokens = int(params.get("max_new_tokens", 1024))
+
+    num_queries = model.config.num_query_tokens
+
+    stop_words = ["Human: ", "Assistant: ", "###", "\n\n"]
+    stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
+    stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+
+    generation_config = GenerationConfig(
+        bos_token_id=1,
+        do_sample=True,
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        stopping_criteria=stopping_criteria
+    )
+
+    pixel_values = None
+    if images is not None:
+        pixel_values = load_image(images).to(device)  # only support one image
+        image_query = DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN
+        prompt = prompt.replace("<image>", image_query)
+
+    elif videos is not None:
+        pixel_values = load_video(videos).to(device)
+        video_query = DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN
+        prompt = prompt.replace("<video>", video_query)
+
+    model_inputs = tokenizer([prompt], return_tensors="pt")
+    model_inputs.pop("token_type_ids", None)
+
+    if pixel_values is not None:
+        model_inputs["pixel_values"] = pixel_values
+
+        generation_output = model.generate(
+            **model_inputs,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+    else:
+        generation_output = model.language_model.generate(
+            **model_inputs,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+
+    preds = generation_output.sequences
+    outputs = tokenizer.batch_decode(preds, skip_special_tokens=True)
+    return outputs
+
+class Chat:
+    def __init__(
+            self,
+            model_path,
+            device,
+            num_gpus=1,
+            load_8bit=False,
+            temperature=0.7,
+            max_new_tokens=512,
+            lora_path=None,
+    ):
+        model, tokenizer = load_model(
+            model_path, device, num_gpus, load_8bit=load_8bit, lora_weights=lora_path
+        )
+
+        self.model = model
+        # self.model.language_model = deepspeed.init_inference(
+        #     self.model.language_model, mp_size=1, dtype=torch.float16, checkpoint=None, replace_with_kernel_inject=True)
+        self.tokenizer = tokenizer
+        num_queries = model.config.num_query_tokens
+
+        self.device = device
+        self.dtype = model.dtype
+
+        stop_words = ["Human: ", "Assistant: ", "###", "\n\n"]
+        stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
+        stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+
+        self.conv = get_conv_template("husky")
+
+        self.image_query = DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN
+        self.video_query = DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN
+
+        self.generation_config = GenerationConfig(
+            bos_token_id=1,
+            do_sample=True,
+            top_k=20,
+            top_p=0.9,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            stopping_criteria=stopping_criteria
+        )
+
+    def ask(self, text, conv, modal_type="image"):
+        assert modal_type in ["text", "image", "video"]
+        conversations = []
+
+        if len(conv.messages) > 0 or modal_type == "text":
+            conv.append_message(conv.roles[0], text)
+        elif modal_type == "image":
+            conv.append_message(conv.roles[0], self.image_query + "\n" + text)
+        else:
+            conv.append_message(conv.roles[0], self.video_query + "\n" + text)
+
+        conv.append_message(conv.roles[1], None)
+        conversations.append(conv.get_prompt())
+        return conversations
+
+    @torch.no_grad()
+    def get_image_embedding(self, image_file):
+        pixel_values = load_image(image_file)
+        pixel_values = pixel_values.unsqueeze(0).to(self.device, dtype=self.dtype)
+        language_model_inputs = self.model.extract_feature(pixel_values)
+        return language_model_inputs
+
+    @torch.no_grad()
+    def get_video_embedding(self, video_file):
+        pixel_values = load_video(video_file)
+        TC, H, W = pixel_values.shape
+        pixel_values = pixel_values.reshape(TC // 3, 3, H, W).transpose(0, 1)  # [C, T, H, W]
+        pixel_values = pixel_values.unsqueeze(0).to(self.device, dtype=self.dtype)
+        assert len(pixel_values.shape) == 5
+        language_model_inputs = self.model.extract_feature(pixel_values)
+        return language_model_inputs
+
+    @torch.no_grad()
+    def answer(self, conversations, language_model_inputs, modal_type="image"):
+        model_inputs = self.tokenizer(
+            conversations,
+            return_tensors="pt",
+        )
+        model_inputs.pop("token_type_ids", None)
+
+        input_ids = model_inputs["input_ids"].to(self.device)
+        attention_mask = model_inputs["attention_mask"].to(self.device)
+
+        if modal_type == "text":
+            generation_output = self.model.language_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                generation_config=self.generation_config,
+                return_dict_in_generate=True,
+                output_scores=True
+            )
+        else:
+            pixel_values = model_inputs.pop("pixel_values", None)
+            if pixel_values is not None:
+                pixel_values = pixel_values.to(self.device)
+
+            generation_output = self.model.generate(
+                pixel_values=pixel_values,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                language_model_inputs=language_model_inputs,
+                generation_config=self.generation_config,
+                return_dict_in_generate=True,
+                output_scores=True
+            )
+
+        preds = generation_output.sequences
+        outputs = self.tokenizer.batch_decode(preds, skip_special_tokens=True)[0]
+
+        if modal_type == "text":
+            skip_echo_len = len(conversations[0]) - conversations[0].count("</s>") * 3
+            outputs = outputs[skip_echo_len:].strip()
+
+        return outputs 
+
+if __name__ == '__main__':
+    # model_path = "/mnt/petrelfs/zhangqinglong/Documents/Husky/work_dirs/husky_v3/EmbodiedGPT/pretrain_0727"
+    model_path = "./"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    chat = Chat(model_path, device=device, num_gpus=1, max_new_tokens=1024, load_8bit=False)
+
+    vision_feature = None
+    image_state = False
+    video_state = False
+
+    while True:
+        query = input("\n")
+        if query.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
+            if os.path.exists(query):
+                print("received.")
+                vision_feature = chat.get_image_embedding(query)
+                chat.conv = get_conv_template("husky").copy()
+                image_state = True
+                continue
+        if query.lower().endswith(('.mp4', '.mkv', '.avi', '.wmv', '.iso', ".webm")):
+            if os.path.exists(query):
+                print("received.")
+                vision_feature = chat.get_video_embedding(query)
+                chat.conv = get_conv_template("husky").copy()
+                video_state = True
+                continue
+
+        if query == "stop":
+            break
+        if query == "clear" or query == "" or query == "\n":
+            chat.conv = get_conv_template("husky").copy()
+            image_state = False
+            video_state = False
+            os.system("clear")
+            print("欢迎使用 husky-13b-zh 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
+            continue
+
+        if image_state:
+            modal_type = "image"
+        elif video_state:
+            modal_type = "video"
+        else:
+            modal_type = "text"
+
+        # image_test = "assets/husky.jpg"
+        # image_test = "assets/yoga.mp4"
+        # video_test = "assets/pretty_girl.mp4"
+        # video_test = "assets/stock-footage-billiards-concentrated-young-woman-playing-in-club.webm"
+        # video_test = "assets/stock-footage-kherson-ukraine-may-open-free-rock-music-festival-crowd-partying-at-a-rock-concert.webm"
+        conversations = chat.ask(text=query, conv=chat.conv, modal_type=modal_type)
+        outputs = chat.answer(conversations, vision_feature, modal_type=modal_type)
+        # NOTE: strip is important to align with the training data.
+        chat.conv.messages[-1][1] = outputs.strip()
+
+        print(f"Husky: \n{outputs}")