import gradio as gr from diffusers import DiffusionPipeline from huggingface_hub import export_folder_as_dduf, create_repo, upload_file import tempfile import torch import os _DTYPE_MAP = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16} article = """ ## DDUF * [DDUF](https://huggingface.co./docs/diffusers/main/en/using-diffusers/other-formats#dduf) Currently, we require a `repo_id` to have all the pipeline components in the Diffusers format. Examples include: [black-forest-labs/FLUX.1-dev](https://huggingface.co./black-forest-labs/FLUX.1-dev), [stabilityai/stable-video-diffusion-img2vid-xt](https://huggingface.co./stabilityai/stable-video-diffusion-img2vid-xt), etc. Partial components will be supported in the future. """ def make_dduf(repo_id: str, destination_repo_id: str, dduf_filename: str, token: str, torch_dtype: str): return_message = "" if destination_repo_id == "": destination_repo_id = repo_id try: destination_repo_id = create_repo(repo_id=destination_repo_id, exist_ok=True).repo_id except Exception as e: return_message += f"❌ Got the following error while creating the repository: \n{e}" with tempfile.TemporaryDirectory() as tmpdir: pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=_DTYPE_MAP[torch_dtype]) if dduf_filename == "": dduf_filename = f"{pipe.__class__.__name__.lower()}" if torch_dtype != "fp32": dduf_filename += f"_{torch_dtype}" dduf_filename += ".dduf" dduf_filename = os.path.join(tmpdir, dduf_filename) pipe.save_pretrained(tmpdir, safe_serialization=True) try: export_folder_as_dduf(dduf_filename, folder_path=tmpdir) except Exception as e: return_message += f"❌ Got the following error while exporting: \n{e}" try: commit_url = upload_file( repo_id=destination_repo_id, path_in_repo=dduf_filename.split("/")[-1], path_or_fileobj=dduf_filename, token=token, ).commit_url return_message += f"Success 🔥. Find the DDUF in [this commit]({commit_url})." except Exception as e: return_message += f"❌ Got the following error while pushing: \n{e}" return str(return_message) demo = gr.Interface( title="DDUF my repo 🤗", article=article, fn=make_dduf, inputs=[ gr.components.Textbox(lines=1, placeholder="Repo ID which should be DDUF'd."), gr.components.Textbox( lines=2, value=None, placeholder="Destination Repo ID that should be used to store the resultant DDUF. Leave it if you want to use the `repo_id` here.", ), gr.components.Textbox( lines=1, value=None, placeholder="Name of the DDUF file. If it's not provided we will infer it based on the pipeline class.", ), gr.components.Textbox(lines=1, placeholder="HF token. You can obtain it from hf.co/settings/tokens."), gr.Dropdown( list(_DTYPE_MAP.keys()), value="fp32", multiselect=False, label="dtype", info="dtype to load the pipeline in.", ), ], outputs="markdown", examples=[ ["stabilityai/stable-video-diffusion-img2vid-xt", "sayakpaul/svd-dduf", "svd.dduf", "hf_XXX", "fp32"], ], allow_flagging="never", ) if __name__ == "__main__": demo.launch(debug=True, show_error=True)