💬 Commit message: Update 2026-01-24 19:24:38, 1 files, 52 lines
📁 Files changed: 1 📝 Lines changed: 52 • sft_get.py
This commit is contained in:
+52
@@ -226,6 +226,44 @@ def display_results(
|
|||||||
_display_civitai_data(civitai_data)
|
_display_civitai_data(civitai_data)
|
||||||
|
|
||||||
|
|
||||||
|
def get_base_name(file_path: Path) -> str:
|
||||||
|
"""Get base filename without .safetensors extension."""
|
||||||
|
name = file_path.name
|
||||||
|
for ext in (".safetensors", ".sft"):
|
||||||
|
if name.lower().endswith(ext):
|
||||||
|
return name[: -len(ext)]
|
||||||
|
return file_path.stem
|
||||||
|
|
||||||
|
|
||||||
|
def save_metadata(
|
||||||
|
file_path: Path,
|
||||||
|
sha256_hash: str,
|
||||||
|
local_metadata: dict[str, Any],
|
||||||
|
civitai_data: dict[str, Any] | None,
|
||||||
|
) -> tuple[Path, Path]:
|
||||||
|
"""Save metadata JSON and SHA256 hash to files alongside the model."""
|
||||||
|
base_name = get_base_name(file_path)
|
||||||
|
parent = file_path.parent
|
||||||
|
|
||||||
|
# Save JSON metadata
|
||||||
|
json_path = parent / f"{base_name}-xm.json"
|
||||||
|
output = {
|
||||||
|
"file": str(file_path),
|
||||||
|
"sha256": sha256_hash,
|
||||||
|
"header_size": local_metadata["header_size"],
|
||||||
|
"tensor_count": local_metadata["tensor_count"],
|
||||||
|
"metadata": local_metadata["metadata"],
|
||||||
|
"civitai": civitai_data,
|
||||||
|
}
|
||||||
|
json_path.write_text(json.dumps(output, indent=2))
|
||||||
|
|
||||||
|
# Save SHA256 hash
|
||||||
|
sha_path = parent / f"{base_name}-xm.sha256"
|
||||||
|
sha_path.write_text(f"{sha256_hash} {file_path.name}\n")
|
||||||
|
|
||||||
|
return json_path, sha_path
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
"""Main entry point."""
|
"""Main entry point."""
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
@@ -254,6 +292,11 @@ def main() -> int:
|
|||||||
dest="json_output",
|
dest="json_output",
|
||||||
help="Output results as JSON",
|
help="Output results as JSON",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--save",
|
||||||
|
action="store_true",
|
||||||
|
help="Save metadata JSON and SHA256 hash alongside the model file",
|
||||||
|
)
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
@@ -292,6 +335,15 @@ def main() -> int:
|
|||||||
else:
|
else:
|
||||||
display_results(file_path, local_metadata, sha256_hash, civitai_data)
|
display_results(file_path, local_metadata, sha256_hash, civitai_data)
|
||||||
|
|
||||||
|
# Save files if requested
|
||||||
|
if args.save:
|
||||||
|
json_path, sha_path = save_metadata(
|
||||||
|
file_path, sha256_hash, local_metadata, civitai_data
|
||||||
|
)
|
||||||
|
console.print()
|
||||||
|
console.print(f"[green]Saved:[/green] {json_path}")
|
||||||
|
console.print(f"[green]Saved:[/green] {sha_path}")
|
||||||
|
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
|
|||||||
Reference in New Issue
Block a user