Skip to content

Commit

Permalink
chore:handle deafault model_id case/model_id being a dict case
Browse files Browse the repository at this point in the history
  • Loading branch information
RUFFY-369 committed Jan 2, 2025
1 parent 7845b6f commit ba02f00
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions runner/app/pipelines/image_to_image_generic.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import logging
import numpy as np
import os
Expand Down Expand Up @@ -45,6 +46,16 @@ def __init__(self, model_id: str, task: str):
kwargs = {"cache_dir": get_model_dir(), "torch_dtype": torch.float16}
torch_device = get_torch_device()

# Check if the model_id is a dictionary in string format in the default value case of model_id on go livepeer side
if model_id.startswith("{") and model_id.endswith("}"):
try:
# Perform json parsing of the string into a dictionary
model_id_dict = json.loads(model_id.replace("'", '"')) # Replace single quotes to make it JSON compliant
if isinstance(model_id_dict, dict) and task in model_id_dict:
model_id = model_id_dict[task]
except json.JSONDecodeError:
pass

folder_name = file_download.repo_folder_name(
repo_id=model_id, repo_type="model"
)
Expand Down

0 comments on commit ba02f00

Please sign in to comment.