import os import sys import requests import json import zipfile import shutil import argparse from tqdm import tqdm import time # Illustrious XL model information MODEL_ID = 795765 MODEL_NAME = "Illustrious XL" MODEL_VERSION = 1 # Version 1.0 MODEL_URL = "https://civitai.com/api/download/models/795765" MODEL_INFO_URL = f"https://civitai.com/api/v1/models/{MODEL_ID}" def download_file(url, destination, filename=None): """Download a file with progress bar""" if filename is None: local_filename = os.path.join(destination, url.split('/')[-1]) else: local_filename = os.path.join(destination, filename) with requests.get(url, stream=True) as r: r.raise_for_status() total_size = int(r.headers.get('content-length', 0)) # Create directory if it doesn't exist os.makedirs(os.path.dirname(local_filename), exist_ok=True) with open(local_filename, 'wb') as f: with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Downloading {os.path.basename(local_filename)}") as pbar: for chunk in r.iter_content(chunk_size=8192): if chunk: f.write(chunk) pbar.update(len(chunk)) return local_filename def create_model_index(model_dir): """Create a model_index.json file for the diffusers library""" model_index = { "_class_name": "StableDiffusionXLPipeline", "_diffusers_version": "0.21.4", "force_zeros_for_empty_prompt": True, "scheduler": [ { "_class_name": "DPMSolverMultistepScheduler", "_diffusers_version": "0.21.4", "beta_end": 0.012, "beta_schedule": "scaled_linear", "beta_start": 0.00085, "num_train_timesteps": 1000, "prediction_type": "epsilon", "solver_order": 2, "solver_type": "midpoint", "thresholding": False, "timestep_spacing": "leading", "trained_betas": None, "use_karras_sigmas": True } ], "text_encoder": [ { "_class_name": "CLIPTextModel", "_diffusers_version": "0.21.4" }, { "_class_name": "CLIPTextModelWithProjection", "_diffusers_version": "0.21.4" } ], "tokenizer": [ { "_class_name": "CLIPTokenizer", "_diffusers_version": "0.21.4" }, { "_class_name": "CLIPTokenizer", "_diffusers_version": "0.21.4" } ], "unet": { "_class_name": "UNet2DConditionModel", "_diffusers_version": "0.21.4" }, "vae": { "_class_name": "AutoencoderKL", "_diffusers_version": "0.21.4" } } with open(os.path.join(model_dir, "model_index.json"), "w") as f: json.dump(model_index, f, indent=2) def download_illustrious_xl(): """Download and set up the Illustrious XL model""" # Set up directories script_dir = os.path.dirname(os.path.abspath(__file__)) models_dir = os.path.join(script_dir, "models") illustrious_dir = os.path.join(models_dir, "illustrious_xl") temp_dir = os.path.join(models_dir, "temp") # Create directories if they don't exist os.makedirs(models_dir, exist_ok=True) os.makedirs(temp_dir, exist_ok=True) # Check if model already exists if os.path.exists(os.path.join(illustrious_dir, "model_index.json")): print(f"⚠️ {MODEL_NAME} model already exists at {illustrious_dir}") choice = input("Do you want to re-download and reinstall the model? (y/n): ") if choice.lower() != 'y': print("Download cancelled.") return # Remove existing model print(f"Removing existing {MODEL_NAME} model...") shutil.rmtree(illustrious_dir, ignore_errors=True) # Create illustrious directory os.makedirs(illustrious_dir, exist_ok=True) # Get model info from Civitai API print(f"Fetching information about {MODEL_NAME} from Civitai...") try: response = requests.get(MODEL_INFO_URL) response.raise_for_status() model_info = response.json() # Save model info for reference with open(os.path.join(illustrious_dir, "model_info.json"), "w") as f: json.dump(model_info, f, indent=2) print(f"Model: {model_info['name']} by {model_info['creator']['username']}") print(f"Description: {model_info['description'][:100]}...") except Exception as e: print(f"⚠️ Failed to fetch model info: {e}") print("Continuing with download anyway...") # Download the model print(f"Downloading {MODEL_NAME} from Civitai...") try: # Download to temp directory model_file = download_file(MODEL_URL, temp_dir, "illustrious_xl.safetensors") # Move the file to the model directory print(f"Setting up {MODEL_NAME} model...") # Create the necessary directory structure for diffusers os.makedirs(os.path.join(illustrious_dir, "unet"), exist_ok=True) os.makedirs(os.path.join(illustrious_dir, "vae"), exist_ok=True) os.makedirs(os.path.join(illustrious_dir, "text_encoder"), exist_ok=True) os.makedirs(os.path.join(illustrious_dir, "text_encoder_2"), exist_ok=True) os.makedirs(os.path.join(illustrious_dir, "tokenizer"), exist_ok=True) os.makedirs(os.path.join(illustrious_dir, "tokenizer_2"), exist_ok=True) # Move the model file to the unet directory shutil.move(model_file, os.path.join(illustrious_dir, "unet", "diffusion_pytorch_model.safetensors")) # Create a model_index.json file create_model_index(illustrious_dir) # Create a README.md file with information about the model with open(os.path.join(illustrious_dir, "README.md"), "w") as f: f.write(f"# {MODEL_NAME}\n\n") f.write(f"Downloaded from Civitai: https://civitai.com/models/{MODEL_ID}\n\n") f.write("This model requires the diffusers library to use.\n") f.write("Use the /generate command in the Discord bot to generate images with this model.\n") print(f"✅ {MODEL_NAME} model has been downloaded and set up successfully!") print(f"Model location: {illustrious_dir}") print("You can now use the model with the /generate command in the Discord bot.") except Exception as e: print(f"❌ Error downloading or setting up the model: {e}") import traceback traceback.print_exc() # Clean up print("Cleaning up...") shutil.rmtree(illustrious_dir, ignore_errors=True) shutil.rmtree(temp_dir, ignore_errors=True) print("Download failed. Please try again later.") return False # Clean up temp directory shutil.rmtree(temp_dir, ignore_errors=True) return True if __name__ == "__main__": parser = argparse.ArgumentParser(description=f"Download and set up the {MODEL_NAME} model from Civitai") parser.add_argument("--force", action="store_true", help="Force download even if the model already exists") args = parser.parse_args() if args.force: # Remove existing model if it exists script_dir = os.path.dirname(os.path.abspath(__file__)) illustrious_dir = os.path.join(script_dir, "models", "illustrious_xl") if os.path.exists(illustrious_dir): print(f"Removing existing {MODEL_NAME} model...") shutil.rmtree(illustrious_dir, ignore_errors=True) download_illustrious_xl()