|
|
import gradio as gr |
|
|
import requests |
|
|
import os |
|
|
import shutil |
|
|
from pathlib import Path |
|
|
from typing import Any |
|
|
from tempfile import TemporaryDirectory |
|
|
from typing import Optional |
|
|
|
|
|
import torch |
|
|
from io import BytesIO |
|
|
|
|
|
from huggingface_hub import CommitInfo, Discussion, HfApi, hf_hub_download |
|
|
from huggingface_hub.file_download import repo_folder_name |
|
|
from diffusers import StableDiffusionXLPipeline |
|
|
from transformers import CONFIG_MAPPING |
|
|
import spaces |
|
|
|
|
|
|
|
|
COMMIT_MESSAGE = " This PR adds fp32 and fp16 weights in safetensors format to {}" |
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
if torch.cuda.is_available(): |
|
|
torch_dtype = torch.bfloat16 |
|
|
else: |
|
|
torch_dtype = torch.float32 |
|
|
|
|
|
@spaces.GPU() |
|
|
def convert_single(model_id: str, filename: str, folder: str, progress: Any, token: str): |
|
|
progress(0, desc="Downloading model") |
|
|
local_file = os.path.join(model_id, filename) |
|
|
ckpt_file = local_file if os.path.isfile(local_file) else hf_hub_download(repo_id=model_id, filename=filename, token=token) |
|
|
|
|
|
pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_file) |
|
|
|
|
|
pipeline.save_pretrained(folder, safe_serialization=True) |
|
|
pipeline = pipeline.to(torch_dtype=torch.float16) |
|
|
pipeline.save_pretrained(folder, safe_serialization=True, variant="fp16") |
|
|
|
|
|
return folder |
|
|
|
|
|
|
|
|
def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: |
|
|
try: |
|
|
discussions = api.get_repo_discussions(repo_id=model_id) |
|
|
except Exception: |
|
|
return None |
|
|
for discussion in discussions: |
|
|
if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title: |
|
|
details = api.get_discussion_details(repo_id=model_id, discussion_num=discussion.num) |
|
|
if details.target_branch == "refs/heads/main": |
|
|
return discussion |
|
|
|
|
|
@spaces.GPU() |
|
|
def convert(token: str, model_id: str, filename: str, progress=gr.Progress()): |
|
|
api = HfApi() |
|
|
|
|
|
pr_title = "Adding `diffusers` weights of this model" |
|
|
|
|
|
with TemporaryDirectory() as d: |
|
|
folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) |
|
|
os.makedirs(folder) |
|
|
new_pr = None |
|
|
try: |
|
|
folder = convert_single(model_id, filename, folder, progress, token) |
|
|
progress(0.7, desc="Uploading to Hub") |
|
|
new_pr = api.upload_folder(folder_path=folder, path_in_repo="./", repo_id=model_id, repo_type="model", token=token, commit_message=pr_title, commit_description=COMMIT_MESSAGE.format(model_id), create_pr=True) |
|
|
pr_number = new_pr.split("%2F")[-1].split("/")[0] |
|
|
link = f"Pr created at: {'https://huggingface.co/' + os.path.join(model_id, 'discussions', pr_number)}" |
|
|
progress(1, desc="Done") |
|
|
except Exception as e: |
|
|
raise gr.exceptions.Error(str(e)) |
|
|
finally: |
|
|
shutil.rmtree(folder) |
|
|
|
|
|
return link |
