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 COMMIT_MESSAGE = " This PR adds fp32 and fp16 weights in safetensors format to {}" 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 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