Skip to content

Commit

Permalink
Fix style
Browse files Browse the repository at this point in the history
  • Loading branch information
mvafin committed Nov 27, 2024
1 parent 33e7deb commit 55a72fc
Showing 1 changed file with 14 additions and 8 deletions.
22 changes: 14 additions & 8 deletions optimum/exporters/openvino/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,13 +338,16 @@ class StoreAttr(object):
path = Path(model_name_or_path)
else:
from diffusers import DiffusionPipeline
path = DiffusionPipeline.download(model_name_or_path,
revision=revision,
cache_dir=cache_dir,
token=token,
local_files_only=local_files_only,
force_download=force_download,
trust_remote_code=trust_remote_code)

path = DiffusionPipeline.download(
model_name_or_path,
revision=revision,
cache_dir=cache_dir,
token=token,
local_files_only=local_files_only,
force_download=force_download,
trust_remote_code=trust_remote_code,
)
model_part_name = None
if (path / "transformer").is_dir():
model_part_name = "transformer"
Expand All @@ -353,12 +356,15 @@ class StoreAttr(object):
dtype = None
if model_part_name:
directory = path / model_part_name
safetensors_files = [filename for filename in directory.glob("*.safetensors") if len(filename.suffixes) == 1]
safetensors_files = [
filename for filename in directory.glob("*.safetensors") if len(filename.suffixes) == 1
]
safetensors_file = None
if len(safetensors_files) > 0:
safetensors_file = safetensors_files.pop(0)
if safetensors_file:
from safetensors import safe_open

with safe_open(safetensors_file, framework="pt", device="cpu") as f:
if len(f.keys()) > 0:
for key in f.keys():
Expand Down

0 comments on commit 55a72fc

Please sign in to comment.