Spaces:
Sleeping
Sleeping
File size: 2,259 Bytes
5372b88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
from pathlib import Path
from huggingface_hub import create_repo, upload_folder
from PIL import Image
def save_model_card(repo_id: str, images=None, base_model=str, dataset_name=str, repo_folder=None):
img_str = ""
for i, image in enumerate(images):
image.save(os.path.join(repo_folder, f"image_{i}.png"))
img_str += f"\n"
yaml = f"""
---
license: creativeml-openrail-m
base_model: {base_model}
tags:
- stable-diffusion
- stable-diffusion-diffusers
- text-to-image
- diffusers
- lora
inference: true
---
"""
model_card = f"""
# LoRA text2image fine-tuning - {repo_id}
These are LoRA adaption weights for {base_model}. The weights were fine-tuned on the {dataset_name} dataset. You can find some example images in the following. \n
{img_str}
"""
with open(os.path.join(repo_folder, "README.md"), "w") as f:
f.write(yaml + model_card)
if __name__ == "__main__":
REPOS = {
"tom_cruise_plain": {"hub_model_id": "person-thumbs-up-plain-lora", "output_dir": "/l/vision/v5/sragas/easel_ai/models_plain/"},
"tom_cruise": {"hub_model_id": "person-thumbs-up-lora", "output_dir": "/l/vision/v5/sragas/easel_ai/models/"},
"tom_cruise_no_cap": {"hub_model_id": "person-thumbs-up-lora-no-cap", "output_dir": "/l/vision/v5/sragas/easel_ai/models_no_cap/"}
}
current_repo_id = "tom_cruise"
current_repo = REPOS[current_repo_id]
print(f"{'-'*20} CURRENT REPO: {current_repo_id} {'-'*20}")
hub_model_id = current_repo["hub_model_id"]
output_dir = current_repo["output_dir"]
pretrained_model_name_or_path = "runwayml/stable-diffusion-v1-5"
images = [Image.open(output_dir+file) for file in os.listdir(output_dir) if ".png" in file]
repo_id = create_repo(
repo_id=hub_model_id or Path(output_dir).name, exist_ok=True, token=None
).repo_id
save_model_card(
repo_id,
images=images,
base_model=pretrained_model_name_or_path,
dataset_name="Custom dataset",
repo_folder=output_dir,
)
upload_folder(
repo_id=repo_id,
folder_path=output_dir,
commit_message="End of training",
ignore_patterns=["step_*", "epoch_*"],
)
|