cleanup in helper

openvinotoolkit · Nov 27, 2024 · 46ef3af · 46ef3af
1 parent 79cd27d
commit 46ef3af
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/notebooks/janus-multimodal-generation/janus-multimodal-generation.ipynb b/notebooks/janus-multimodal-generation/janus-multimodal-generation.ipynb
@@ -78,9 +78,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%pip install -q \"gradio>=4.19\" \"torch>=2.2\" \"safetensors\" \"transformers>=4.38\" \"nncf>=2.14\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
-    "%pip install -q \"git+https://github.com/deepseek-ai/Janus\"\n",
-    "%pip install -U --pre \"openvino>2024.5\" --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly"
+    "import platform\n",
+    "\n",
+    "%pip install -q \"gradio>=4.19\" \"torch>=2.2\" \"torchvision\" \"safetensors\" \"transformers>=4.38\" \"nncf>=2.14\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
+    "%pip install -q \"git+https://github.com/deepseek-ai/Janus\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
+    "%pip install -U --pre \"openvino>2024.5\" --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly\n",
+    "\n",
+    "if platform.system() == \"Darwin\":\n",
+    "    %pip install -q \"numpy<2.0.0\""
    ]
   },
   {

diff --git a/notebooks/janus-multimodal-generation/ov_janus_helper.py b/notebooks/janus-multimodal-generation/ov_janus_helper.py
@@ -1,17 +1,17 @@
 from pathlib import Path
 import types
-from typing import Optional, Tuple, Union, List
+from typing import Optional, Tuple, List
 import gc
 import openvino as ov
 from openvino.runtime import opset13
 import nncf
 import numpy as np
 from tqdm.auto import tqdm
 import torch
-from transformers import AutoModelForCausalLM, AutoProcessor, AutoConfig, TextStreamer
+from transformers import AutoModelForCausalLM, AutoConfig
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from transformers.generation import GenerationConfig, GenerationMixin
-from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from transformers.modeling_outputs import CausalLMOutputWithPast
 from janus.utils.io import load_pil_images
 from PIL import Image
 
@@ -231,7 +231,10 @@ def convert_janus_model(model_id, output_dir, quantization_config):
     print(f"⌛ {model_name} conversion started. Be patient, it may takes some time.")
     print("⌛ Load Original model")
     processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_id)
-    vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+    config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    language_config = config.language_config
+    language_config._attn_implementation = 'sdpa'
+    vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_id, language_config=language_config, trust_remote_code=True)
     vl_gpt = vl_gpt.eval()
     vl_gpt.config.save_pretrained(output_dir)
     processor.save_pretrained(output_dir)