Prevent unwanted log messages from modules

extensions/multimodal/multimodal_embedder.py

@@ -1,5 +1,4 @@
 import base64
-import logging
 import re
 from dataclasses import dataclass
 from io import BytesIO
@@ -10,6 +9,7 @@ from PIL import Image
 
 from extensions.multimodal.pipeline_loader import load_pipeline
 from modules import shared
+from modules.logging_colors import logger
 from modules.text_generation import encode, get_max_prompt_length
 
 
@@ -26,7 +26,7 @@ class MultimodalEmbedder:
     def __init__(self, params: dict):
         pipeline, source = load_pipeline(params)
         self.pipeline = pipeline
-        logging.info(f'Multimodal: loaded pipeline {self.pipeline.name()} from pipelines/{source} ({self.pipeline.__class__.__name__})')
+        logger.info(f'Multimodal: loaded pipeline {self.pipeline.name()} from pipelines/{source} ({self.pipeline.__class__.__name__})')
 
     def _split_prompt(self, prompt: str, load_images: bool = False) -> List[PromptPart]:
         """Splits a prompt into a list of `PromptParts` to separate image data from text.
@@ -138,7 +138,7 @@ class MultimodalEmbedder:
 
         # notify user if we truncated an image
         if removed_images > 0:
-            logging.warning(f"Multimodal: removed {removed_images} image(s) from prompt. Try decreasing max_new_tokens if generation is broken")
+            logger.warning(f"Multimodal: removed {removed_images} image(s) from prompt. Try decreasing max_new_tokens if generation is broken")
 
         return encoded
 

extensions/multimodal/pipeline_loader.py

@@ -1,4 +1,3 @@
-import logging
 import traceback
 from importlib import import_module
 from pathlib import Path
@@ -6,6 +5,7 @@ from typing import Tuple
 
 from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
 from modules import shared
+from modules.logging_colors import logger
 
 
 def _get_available_pipeline_modules():
@@ -21,8 +21,8 @@ def load_pipeline(params: dict) -> Tuple[AbstractMultimodalPipeline, str]:
         try:
             pipeline_modules[name] = import_module(f'extensions.multimodal.pipelines.{name}.pipelines')
         except:
-            logging.warning(f'Failed to get multimodal pipelines from {name}')
-            logging.warning(traceback.format_exc())
+            logger.warning(f'Failed to get multimodal pipelines from {name}')
+            logger.warning(traceback.format_exc())
 
     if shared.args.multimodal_pipeline is not None:
         for k in pipeline_modules:
@@ -48,5 +48,5 @@ def load_pipeline(params: dict) -> Tuple[AbstractMultimodalPipeline, str]:
         log = f'Multimodal - ERROR: Failed to load multimodal pipeline "{shared.args.multimodal_pipeline}", available pipelines are: {available}.'
     else:
         log = f'Multimodal - ERROR: Failed to determine multimodal pipeline for model {shared.args.model}, please select one manually using --multimodal-pipeline [PIPELINE]. Available pipelines are: {available}.'
-    logging.critical(f'{log} Please specify a correct pipeline, or disable the extension')
+    logger.critical(f'{log} Please specify a correct pipeline, or disable the extension')
     raise RuntimeError(f'{log} Please specify a correct pipeline, or disable the extension')

extensions/multimodal/pipelines/llava/llava.py

@@ -1,16 +1,17 @@
-import logging
 import time
 from abc import abstractmethod
 from typing import List, Tuple
 
 import torch
-from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
 from huggingface_hub import hf_hub_download
-from modules import shared
-from modules.text_generation import encode
 from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModel
 
+from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
+from modules import shared
+from modules.logging_colors import logger
+from modules.text_generation import encode
+
 
 class LLaVA_v0_Pipeline(AbstractMultimodalPipeline):
     CLIP_REPO = "openai/clip-vit-large-patch14"
@@ -26,11 +27,11 @@ class LLaVA_v0_Pipeline(AbstractMultimodalPipeline):
     def _load_models(self):
         start_ts = time.time()
 
-        logging.info(f"LLaVA - Loading CLIP from {LLaVA_v0_Pipeline.CLIP_REPO} as {self.clip_dtype} on {self.clip_device}...")
+        logger.info(f"LLaVA - Loading CLIP from {LLaVA_v0_Pipeline.CLIP_REPO} as {self.clip_dtype} on {self.clip_device}...")
         image_processor = CLIPImageProcessor.from_pretrained(LLaVA_v0_Pipeline.CLIP_REPO, torch_dtype=self.clip_dtype)
         vision_tower = CLIPVisionModel.from_pretrained(LLaVA_v0_Pipeline.CLIP_REPO, torch_dtype=self.clip_dtype).to(self.clip_device)
 
-        logging.info(f"LLaVA - Loading projector from {self.llava_projector_repo()} as {self.projector_dtype} on {self.projector_device}...")
+        logger.info(f"LLaVA - Loading projector from {self.llava_projector_repo()} as {self.projector_dtype} on {self.projector_device}...")
         projector_path = hf_hub_download(self.llava_projector_repo(), self.llava_projector_filename())
         mm_projector = torch.nn.Linear(*self.llava_projector_shape())
         projector_data = torch.load(projector_path)
@@ -38,7 +39,7 @@ class LLaVA_v0_Pipeline(AbstractMultimodalPipeline):
         mm_projector.bias = torch.nn.Parameter(projector_data['model.mm_projector.bias'].to(dtype=self.projector_dtype), False)
         mm_projector = mm_projector.to(self.projector_device)
 
-        logging.info(f"LLaVA supporting models loaded, took {time.time() - start_ts:.2f} seconds")
+        logger.info(f"LLaVA supporting models loaded, took {time.time() - start_ts:.2f} seconds")
         return image_processor, vision_tower, mm_projector
 
     @staticmethod

extensions/multimodal/script.py

@@ -1,5 +1,4 @@
 import base64
-import logging
 import re
 import time
 from functools import partial
@@ -10,6 +9,7 @@ import torch
 
 from extensions.multimodal.multimodal_embedder import MultimodalEmbedder
 from modules import shared
+from modules.logging_colors import logger
 
 params = {
     "add_all_images_to_prompt": False,
@@ -78,7 +78,7 @@ def tokenizer_modifier(state, prompt, input_ids, input_embeds):
         return prompt, input_ids, input_embeds
 
     prompt, input_ids, input_embeds, total_embedded = multimodal_embedder.forward(prompt, state, params)
-    logging.info(f'Embedded {total_embedded} image(s) in {time.time()-start_ts:.2f}s')
+    logger.info(f'Embedded {total_embedded} image(s) in {time.time()-start_ts:.2f}s')
     return (prompt,
             input_ids.unsqueeze(0).to(shared.model.device, dtype=torch.int64),
             input_embeds.unsqueeze(0).to(shared.model.device, dtype=shared.model.dtype))