    def predict(
        self,
        image: Path = Input(description="Input image for segmentation"),
        instruction: str = Input(description="Text instruction for the model"),
    ) -> Output:
        """Run a single prediction on the model"""
        # Prepare the image
        image = Image.open(str(image)).convert('RGB')
        
        # Prepare the input
        text_prompts = f"<image>{instruction}"
        input_dict = {
            'image': image,
            'text': text_prompts,
            'past_text': '',
            'mask_prompts': None,
            'tokenizer': self.tokenizer,
        }
        
        # Get model prediction
        return_dict = self.model.predict_forward(**input_dict)
        answer = return_dict["prediction"]
        
        # Handle segmentation if present
        output_path = None
        if '[SEG]' in answer:
            pred_masks = return_dict["prediction_masks"][0]
            
            # Ensure mask is in the correct format
            if isinstance(pred_masks, np.ndarray):
                binary_mask = (pred_masks > 0.5).astype('uint8') * 255
            else:
                binary_mask = (pred_masks.cpu().numpy() > 0.5).astype('uint8') * 255
            
            # Ensure mask has valid dimensions
            if binary_mask.ndim == 2:
                height, width = binary_mask.shape
            elif binary_mask.ndim == 3:
                # If we have a 3D array, take the first channel
                binary_mask = binary_mask[0] if binary_mask.shape[0] == 1 else binary_mask[:, :, 0]
                height, width = binary_mask.shape
            else:
                return Output(img=None, response=str(answer))
                
            # Check if dimensions are valid and mask is not empty
            if width > 0 and height > 0 and np.any(binary_mask):
                # Create output directory if it doesn't exist
                os.makedirs("/tmp", exist_ok=True)
                
                # Save the binary mask
                output_path = "/tmp/output.png"
                if cv2.imwrite(output_path, binary_mask):
                    return Output(img=Path(output_path), response=str(answer))

        return Output(img=None, response=str(answer))
Run a single prediction on the model
demo/predict-img.py