import os import gradio as gr import torch from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTextConfig, CLIPTokenizer from safetensors.torch import load_file from collections import OrderedDict import re import json import requests import subprocess from urllib.parse import urlparse, unquote from pathlib import Path import hashlib from datetime import datetime from typing import Dict, List, Optional from huggingface_hub import login, HfApi, hf_hub_download from huggingface_hub.utils import validate_repo_id, HFValidationError from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE from huggingface_hub.utils import HfHubHTTPError # ---------------------- DEPENDENCIES ---------------------- def install_dependencies_gradio(): """Installs the necessary dependencies.""" try: subprocess.run( [ "pip", "install", "-U", "torch", "diffusers", "transformers", "accelerate", "safetensors", "huggingface_hub", "xformers", ] ) print("Dependencies installed successfully.") except Exception as e: print(f"Error installing dependencies: {e}") # ---------------------- UTILITY FUNCTIONS ---------------------- def increment_filename(filename): """Increments the filename to avoid overwriting existing files.""" base, ext = os.path.splitext(filename) counter = 1 while os.path.exists(filename): filename = f"{base}({counter}){ext}" counter += 1 return filename # ---------------------- UPLOAD FUNCTION ---------------------- def create_model_repo(api, user, orgs_name, model_name, make_private=False): """Creates a Hugging Face model repository.""" repo_id = ( f"{orgs_name}/{model_name.strip()}" if orgs_name else f"{user['name']}/{model_name.strip()}" ) try: api.create_repo(repo_id=repo_id, repo_type="model", private=make_private) print(f"Model repo '{repo_id}' created.") except HfHubHTTPError: print(f"Model repo '{repo_id}' already exists.") return repo_id # ---------------------- MODEL LOADING AND CONVERSION ---------------------- def download_model(model_path_or_url): """Downloads a model, handling URLs, HF repos, and local paths.""" try: # 1. Check if it's a valid Hugging Face repo ID try: validate_repo_id(model_path_or_url) local_path = hf_hub_download(repo_id=model_path_or_url) return local_path except HFValidationError: pass # 2. Check if it's a URL if model_path_or_url.startswith("http://") or model_path_or_url.startswith("https://"): response = requests.get(model_path_or_url, stream=True) response.raise_for_status() parsed_url = urlparse(model_path_or_url) filename = os.path.basename(unquote(parsed_url.path)) if not filename: filename = hashlib.sha256(model_path_or_url.encode()).hexdigest() cache_dir = os.path.join(HUGGINGFACE_HUB_CACHE, "downloads") os.makedirs(cache_dir, exist_ok=True) local_path = os.path.join(cache_dir, filename) with open(local_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) return local_path # 3. Check if it's a local file elif os.path.isfile(model_path_or_url): return model_path_or_url # 4. Handle Hugging Face repo with a specific file else: try: parts = model_path_or_url.split("/", 1) if len(parts) == 2: repo_id, filename = parts validate_repo_id(repo_id) local_path = hf_hub_download(repo_id=repo_id, filename=filename) return local_path else: raise ValueError("Invalid input format.") except HFValidationError: raise ValueError(f"Invalid model path or URL: {model_path_or_url}") except Exception as e: raise ValueError(f"Error downloading or accessing model: {e}") def load_sdxl_checkpoint(checkpoint_path): """Loads checkpoint and extracts state dicts, handling Illustrious-xl.""" if checkpoint_path.endswith(".safetensors"): state_dict = load_file(checkpoint_path, device="cpu") elif checkpoint_path.endswith(".ckpt"): state_dict = torch.load(checkpoint_path, map_location="cpu")["state_dict"] else: raise ValueError("Unsupported checkpoint format. Must be .safetensors or .ckpt") text_encoder1_state = OrderedDict() text_encoder2_state = OrderedDict() vae_state = OrderedDict() unet_state = OrderedDict() for key, value in state_dict.items(): if key.startswith("first_stage_model."): # VAE vae_state[key.replace("first_stage_model.", "")] = value.to(torch.float16) elif key.startswith("condition_model.model.text_encoder."): # First Text Encoder text_encoder1_state[key.replace("condition_model.model.text_encoder.", "")] = value.to(torch.float16) elif key.startswith("condition_model.model.text_encoder_2."): # Second Text Encoder text_encoder2_state[key.replace("condition_model.model.text_encoder_2.", "")] = value.to(torch.float16) elif key.startswith("model.diffusion_model."): # UNet unet_state[key.replace("model.diffusion_model.", "")] = value.to(torch.float16) return text_encoder1_state, text_encoder2_state, vae_state, unet_state def build_diffusers_model( text_encoder1_state, text_encoder2_state, vae_state, unet_state, reference_model_path=None ): """Builds Diffusers components, loading state dicts with strict=False.""" if not reference_model_path: reference_model_path = "stabilityai/stable-diffusion-xl-base-1.0" # Load configurations from the reference model config_text_encoder1 = CLIPTextConfig.from_pretrained( reference_model_path, subfolder="text_encoder" ) config_text_encoder2 = CLIPTextConfig.from_pretrained( reference_model_path, subfolder="text_encoder_2" ) config_vae = AutoencoderKL.from_pretrained(reference_model_path, subfolder="vae").config config_unet = UNet2DConditionModel.from_pretrained(reference_model_path, subfolder="unet").config # Create instances using the configurations text_encoder1 = CLIPTextModel(config_text_encoder1) text_encoder2 = CLIPTextModelWithProjection(config_text_encoder2) # Use CLIPTextModelWithProjection vae = AutoencoderKL(config=config_vae) unet = UNet2DConditionModel(config=config_unet) # Load state dicts with strict=False text_encoder1.load_state_dict(text_encoder1_state, strict=False) text_encoder2.load_state_dict(text_encoder2_state, strict=False) vae.load_state_dict(vae_state, strict=False) unet.load_state_dict(unet_state, strict=False) text_encoder1.to(torch.float16).to("cpu") text_encoder2.to(torch.float16).to("cpu") vae.to(torch.float16).to("cpu") unet.to(torch.float16).to("cpu") return text_encoder1, text_encoder2, vae, unet def convert_and_save_sdxl_to_diffusers( checkpoint_path_or_url, output_path, reference_model_path ): """Converts and saves the Illustrious-xl checkpoint to Diffusers format.""" checkpoint_path = download_model(checkpoint_path_or_url) text_encoder1_state, text_encoder2_state, vae_state, unet_state = ( load_sdxl_checkpoint(checkpoint_path) ) text_encoder1, text_encoder2, vae, unet = build_diffusers_model( text_encoder1_state, text_encoder2_state, vae_state, unet_state, reference_model_path, ) # Load tokenizer and scheduler from the reference model pipeline = StableDiffusionXLPipeline.from_pretrained( reference_model_path, text_encoder=text_encoder1, text_encoder_2=text_encoder2, vae=vae, unet=unet, torch_dtype=torch.float16, ) pipeline.to("cpu") pipeline.save_pretrained(output_path) print(f"Model saved as Diffusers format: {output_path}") # ---------------------- UPLOAD FUNCTION ---------------------- def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private): """Uploads a model to the Hugging Face Hub.""" login(token=hf_token, add_to_git_credential=True) api = HfApi() user = api.whoami(token=hf_token) model_repo = create_model_repo(api, user, orgs_name, model_name, make_private) api.upload_folder(folder_path=model_path, repo_id=model_repo) print(f"Model uploaded to: https://huggingface.co/{model_repo}") # ---------------------- GRADIO INTERFACE ---------------------- def main( model_to_load, reference_model, output_path, hf_token, orgs_name, model_name, make_private, ): """Main function: SDXL checkpoint to Diffusers, always fp16.""" try: convert_and_save_sdxl_to_diffusers( model_to_load, output_path, reference_model ) upload_to_huggingface( output_path, hf_token, orgs_name, model_name, make_private ) return "Conversion and upload completed successfully!" except Exception as e: return f"An error occurred: {e}" # Return the error message css = """ #main-container { display: flex; flex-direction: column; font-family: 'Arial', sans-serif; font-size: 16px; color: #333; } #convert-button { margin-top: 1em; /* Adds some space above the button */ } """ with gr.Blocks(css=css) as demo: gr.Markdown( """ # 🎨 SDXL Model Converter Convert SDXL checkpoints to Diffusers format (FP16, CPU-only). ### đŸ“Ĩ Input Sources Supported: - Local model files (.safetensors, .ckpt) - Direct URLs to model files - Hugging Face model repositories (e.g., 'my-org/my-model' or 'my-org/my-model/file.safetensors') ### ℹī¸ Important Notes: - This tool runs on **CPU**, conversion might be slower than on GPU. - For Hugging Face uploads, you need a **WRITE** token (not a read token). - Get your HF token here: [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) ### 💾 Memory Usage: - This space is configured for **FP16** precision to reduce memory usage. - Close other applications during conversion. - For large models, ensure you have at least 16GB of RAM. ### đŸ’ģ Source Code: - [GitHub Repository](https://github.com/Ktiseos-Nyx/Gradio-SDXL-Diffusers) ### 🙏 Support: - If you're interested in funding more projects: [Ko-fi](https://ko-fi.com/duskfallcrew) """ ) with gr.Row(): with gr.Column(): model_to_load = gr.Textbox( label="SDXL Checkpoint (Path, URL, or HF Repo)", placeholder="Path, URL, or Hugging Face Repo ID (e.g., my-org/my-model or my-org/my-model/file.safetensors)", ) reference_model = gr.Textbox( label="Reference Diffusers Model (Optional)", placeholder="e.g., stabilityai/stable-diffusion-xl-base-1.0 (Leave blank for default)", ) output_path = gr.Textbox( label="Output Path (Diffusers Format)", value="output" ) hf_token = gr.Textbox( label="Hugging Face Token", placeholder="Your Hugging Face write token", type="password" ) orgs_name = gr.Textbox( label="Organization Name (Optional)", placeholder="Your organization name" ) model_name = gr.Textbox( label="Model Name", placeholder="The name of your model on Hugging Face" ) make_private = gr.Checkbox(label="Make Repository Private", value=False) convert_button = gr.Button("Convert and Upload") with gr.Column(variant="panel"): # Use variant="panel" output = gr.Markdown(container=False) convert_button.click( fn=main, inputs=[ model_to_load, reference_model, output_path, hf_token, orgs_name, model_name, make_private, ], outputs=output, ) demo.launch()