59 lines
1.9 KiB
Python
Executable File
59 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import replicate
|
|
import json
|
|
import os
|
|
import argparse
|
|
|
|
|
|
def format_json_file(file_path):
|
|
try:
|
|
with open(file_path, "r") as f:
|
|
data = json.load(f)
|
|
data["run_count"] = 0
|
|
|
|
with open(file_path, "w") as f:
|
|
json.dump(data, f, indent=4, ensure_ascii=False)
|
|
except json.JSONDecodeError:
|
|
print(f"Error: {file_path} contains invalid JSON")
|
|
except IOError:
|
|
print(f"Error: Could not read or write to {file_path}")
|
|
|
|
|
|
def format_json_files_in_directory(directory):
|
|
for filename in os.listdir(directory):
|
|
if filename.endswith(".json"):
|
|
file_path = os.path.join(directory, filename)
|
|
format_json_file(file_path)
|
|
|
|
|
|
def update_schemas(update=False):
|
|
with open("supported_models.json", "r", encoding="utf-8") as f:
|
|
supported_models = json.load(f)
|
|
|
|
schemas_directory = "schemas"
|
|
existing_schemas = set(os.listdir(schemas_directory))
|
|
|
|
for model in supported_models["models"]:
|
|
schema_filename = f"{model.replace('/', '_')}.json"
|
|
schema_path = os.path.join(schemas_directory, schema_filename)
|
|
|
|
if update or schema_filename not in existing_schemas:
|
|
try:
|
|
m = replicate.models.get(model)
|
|
with open(schema_path, "w", encoding="utf-8") as f:
|
|
f.write(m.json())
|
|
print(f"{'Updated' if update else 'Added'} schema for {model}")
|
|
except replicate.exceptions.ReplicateError as e:
|
|
print(f"Error fetching schema for {model}: {str(e)}")
|
|
continue
|
|
|
|
format_json_files_in_directory(schemas_directory)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Update model schemas")
|
|
parser.add_argument("--update", action="store_true", help="Update all schemas, not just new ones")
|
|
args = parser.parse_args()
|
|
|
|
update_schemas(update=args.update)
|