Upload folder using huggingface_hub

geopixel.py

@@ -12,6 +12,7 @@ from model.IXC.modeling_internlm_xcomposer2 import InternLMXComposer2ForCausalLM
 from model.IXC.modeling_internlm2 import InternLM2Model
 from model.sam2.build_sam import build_sam2_hf
 from model.sam2.utils.transforms import SAM2Transforms
+from transformers import TextStreamer
 try:
     from transformers.generation.streamers import BaseStreamer
 except:  # noqa # pylint: disable=bare-except
@@ -93,8 +94,10 @@ class GeoPixelMetaModel:
             (128, 128),
             (64, 64),
         ]
+        
         for param in self.visual_model.parameters():
             param.requires_grad = False
+
         if config.train_mask_decoder:
             self.visual_model.sam_mask_decoder.train()
             for param in self.visual_model.sam_mask_decoder.parameters():
@@ -195,6 +198,8 @@ class GeoPixelForCausalLM(InternLMXComposer2ForCausalLM):
         samples = kwargs.get('samples', None)
         if samples and samples['data_type'][0] == 'grounding': 
             kwargs['output_hidden_states'] = True
+            kwargs['use_cache'] = False
+
             torch.cuda.empty_cache()
             outputs = super().forward(**kwargs)
 
@@ -246,9 +251,6 @@ class GeoPixelForCausalLM(InternLMXComposer2ForCausalLM):
                     low_res_masks,
                     ori_hw[i],
                 )
-
-                # pred_masks = pred_masks.squeeze(0) 
-                # all_pred_masks.append(pred_masks)
                 all_pred_masks.append(pred_masks[:, 0])
                 
 
@@ -320,27 +322,32 @@ class GeoPixelForCausalLM(InternLMXComposer2ForCausalLM):
         hd_num: int = 9,
         history: List[Tuple[str, str]] = [],
         max_new_tokens: int = 1024,
+        stream: bool = False,
         **kwargs,
     ):
         with torch.no_grad():
             inputs, im_mask, _ = self.interleav_wrap_chat(query, images, history=history, hd_num=hd_num)        
-            print(im_mask.sum().item())
             inputs = {
                 k: v.to(self.device)
                 for k, v in inputs.items() if torch.is_tensor(v)
             }
-            # print(len(inputs['inputs_embeds'][0]))
             eos_token_id = [
                 tokenizer.eos_token_id,
                 #tokenizer.convert_tokens_to_ids(['[UNUSED_TOKEN_145]'])[0]
             ]
             all_pred_masks = []
+            
+            if stream:
+                streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            else: 
+                streamer = None
+
             outputs = self.generate(
                 **inputs,
                 max_new_tokens=max_new_tokens,
                 im_mask=im_mask,
                 input_ids = None,
-                streamer= None,
+                streamer= streamer,
                 num_beams=1,
                 do_sample=False,
                 temperature=1.0,