Update app.py

UGH THE CONVERSION AND DICT WAS NOT IN THERE BEFORE LAZY LLM

app.py

@@ -1,35 +1,68 @@
 import os
 import gradio as gr
 import torch
-from diffusers import StableDiffusionXLPipeline
-from huggingface_hub import HfApi, login
-from huggingface_hub.utils import validate_repo_id, HfHubHTTPError
-import tempfile
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTextConfig
+from safetensors.torch import load_file
+from collections import OrderedDict
 import re
 import json
-import glob
 import gdown
 import requests
 import subprocess
 from urllib.parse import urlparse, unquote
 from pathlib import Path
+import tempfile
+from tqdm import tqdm
 
 # ---------------------- UTILITY FUNCTIONS ----------------------
 
-def get_save_dtype(save_precision_as):
-    """Determines the save dtype based on the user's choice."""
-    if save_precision_as == "fp16":
-        return torch.float16
-    elif save_precision_as == "bf16":
-        return torch.bfloat16
-    elif save_precision_as == "float":
-        return torch.float32  # Using float32 for "float" option
+def is_valid_url(url):
+    """Checks if a string is a valid URL."""
+    try:
+        result = urlparse(url)
+        return all([result.scheme, result.netloc])
+    except:
+        return False
+
+def get_filename(url):
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    if 'content-disposition' in response.headers:
+        content_disposition = response.headers['content-disposition']
+        filename = re.findall('filename="?([^"]+)"?', content_disposition)[0]
     else:
-        return None
+        url_path = urlparse(url).path
+        filename = unquote(os.path.basename(url_path))
+
+    return filename
+
+def get_supported_extensions():
+    return tuple([".ckpt", ".safetensors", ".pt", ".pth"])
+
+def download_model(url, dst, output_widget):
+    filename = get_filename(url)
+    filepath = os.path.join(dst, filename)
+    try:
+        if "drive.google.com" in url:
+            gdown = gdown_download(url, dst, filepath)
+        else:
+            if "huggingface.co" in url:
+                if "/blob/" in url:
+                    url = url.replace("/blob/", "/resolve/")
+            subprocess.run(["aria2c","-x 16",url,"-d",dst,"-o",filename])
+        with output_widget:
+            return filepath
+    except Exception as e:
+       with output_widget:
+            return None
 
 def determine_load_checkpoint(model_to_load):
-    """Determines if the model to load is a checkpoint or a Diffusers model."""
-    if model_to_load.endswith('.ckpt') or model_to_load.endswith('.safetensors'):
+    """Determines if the model to load is a checkpoint, Diffusers model, or URL."""
+    if is_valid_url(model_to_load) and (model_to_load.endswith(get_supported_extensions())):
+        return True
+    elif model_to_load.endswith(get_supported_extensions()):
         return True
     elif os.path.isdir(model_to_load):
         required_folders = {"unet", "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler", "vae"}
@@ -37,15 +70,6 @@ def determine_load_checkpoint(model_to_load):
             return False
     return None  # handle this case as required
 
-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
-
 def create_model_repo(api, user, orgs_name, model_name, make_private=False):
     """Creates a Hugging Face model repository if it doesn't exist."""
     if orgs_name == "":
@@ -69,6 +93,112 @@ def is_diffusers_model(model_path):
     required_folders = {"unet", "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler", "vae"}
     return required_folders.issubset(set(os.listdir(model_path))) and os.path.isfile(os.path.join(model_path, "model_index.json"))
 
+# ---------------------- MODEL UTIL (From library.sdxl_model_util) ----------------------
+def load_models_from_sdxl_checkpoint(sdxl_base_id, checkpoint_path, device):
+    """Loads SDXL model components from a checkpoint file."""
+    text_encoder1 = CLIPTextModel.from_pretrained(sdxl_base_id, subfolder="text_encoder").to(device)
+    text_encoder2 = CLIPTextModel.from_pretrained(sdxl_base_id, subfolder="text_encoder_2").to(device)
+    vae = AutoencoderKL.from_pretrained(sdxl_base_id, subfolder="vae").to(device)
+    unet = UNet2DConditionModel.from_pretrained(sdxl_base_id, subfolder="unet").to(device)
+    unet = unet
+
+    ckpt_state_dict = torch.load(checkpoint_path, map_location=device)
+
+    o = OrderedDict()
+    for key in list(ckpt_state_dict.keys()):
+        o[key.replace("module.", "")] = ckpt_state_dict[key]
+    del ckpt_state_dict
+
+    print("Applying weights to text encoder 1:")
+    text_encoder1.load_state_dict({
+        '.'.join(key.split('.')[1:]): o[key] for key in list(o.keys()) if key.startswith("first_stage_model.cond_stage_model.model.transformer")
+    }, strict=False)
+    print("Applying weights to text encoder 2:")
+    text_encoder2.load_state_dict({
+        '.'.join(key.split('.')[1:]): o[key] for key in list(o.keys()) if key.startswith("cond_stage_model.model.transformer")
+    }, strict=False)
+    print("Applying weights to VAE:")
+    vae.load_state_dict({
+        '.'.join(key.split('.')[2:]): o[key] for key in list(o.keys()) if key.startswith("first_stage_model.model")
+    }, strict=False)
+    print("Applying weights to UNet:")
+    unet.load_state_dict({
+        key: o[key] for key in list(o.keys()) if key.startswith("model.diffusion_model")
+    }, strict=False)
+
+    logit_scale = None #Not used here!
+    global_step = None #Not used here!
+    return text_encoder1, text_encoder2, vae, unet, logit_scale, global_step
+
+def save_stable_diffusion_checkpoint(save_path, text_encoder1, text_encoder2, unet, epoch, global_step, ckpt_info, vae, logit_scale, save_dtype):
+    """Saves the stable diffusion checkpoint."""
+    weights = OrderedDict()
+    text_encoder1_dict = text_encoder1.state_dict()
+    text_encoder2_dict = text_encoder2.state_dict()
+    unet_dict = unet.state_dict()
+    vae_dict = vae.state_dict()
+
+    def replace_key(key):
+        key = "cond_stage_model.model.transformer." + key
+        return key
+
+    print("Merging text encoder 1")
+    for key in tqdm(list(text_encoder1_dict.keys())):
+        weights["first_stage_model.cond_stage_model.model.transformer." + key] = text_encoder1_dict[key].to(save_dtype)
+
+    print("Merging text encoder 2")
+    for key in tqdm(list(text_encoder2_dict.keys())):
+        weights[replace_key(key)] = text_encoder2_dict[key].to(save_dtype)
+
+    print("Merging vae")
+    for key in tqdm(list(vae_dict.keys())):
+        weights["first_stage_model.model." + key] = vae_dict[key].to(save_dtype)
+
+    print("Merging unet")
+    for key in tqdm(list(unet_dict.keys())):
+        weights["model.diffusion_model." + key] = unet_dict[key].to(save_dtype)
+
+    info = {"epoch": epoch, "global_step": global_step}
+    if ckpt_info is not None:
+        info.update(ckpt_info)
+
+    if logit_scale is not None:
+        info["logit_scale"] = logit_scale.item()
+
+    torch.save({"state_dict": weights, "info": info}, save_path)
+
+    key_count = len(weights.keys())
+    del weights
+    del text_encoder1_dict, text_encoder2_dict, unet_dict, vae_dict
+    return key_count
+
+def save_diffusers_checkpoint(save_path, text_encoder1, text_encoder2, unet, reference_model, vae, trim_if_model_exists, save_dtype):
+    """Saves Diffusers-style checkpoint from the model."""
+    print("Saving SDXL as Diffusers format to:", save_path)
+    print("SDXL Text Encoder 1 to:", os.path.join(save_path, "text_encoder"))
+    text_encoder1.save_pretrained(os.path.join(save_path, "text_encoder"))
+
+    print("SDXL Text Encoder 2 to:", os.path.join(save_path, "text_encoder_2"))
+    text_encoder2.save_pretrained(os.path.join(save_path, "text_encoder_2"))
+
+    print("SDXL VAE to:", os.path.join(save_path, "vae"))
+    vae.save_pretrained(os.path.join(save_path, "vae"))
+
+    print("SDXL UNet to:", os.path.join(save_path, "unet"))
+    unet.save_pretrained(os.path.join(save_path, "unet"))
+
+    if reference_model is not None:
+        print(f"Copying scheduler from {reference_model}")
+        scheduler_src = StableDiffusionXLPipeline.from_pretrained(reference_model, torch_dtype=torch.float16).scheduler
+        torch.save(scheduler_src.config, os.path.join(save_path, "scheduler", "scheduler_config.json"))
+    else:
+        print(f"No reference Model. Copying scheduler from original model.")
+        scheduler_src = StableDiffusionXLPipeline.from_pretrained(reference_model, torch_dtype=torch.float16).scheduler
+        scheduler_src.save_pretrained(save_path)
+
+    if trim_if_model_exists:
+        print("Trim Complete")
+
 # ---------------------- CONVERSION AND UPLOAD FUNCTIONS ----------------------
 
 def load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget):
@@ -86,15 +216,43 @@ def load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget):
 
 def load_from_sdxl_checkpoint(args, output_widget):
     """Loads the SDXL model components from a checkpoint file (placeholder)."""
-    # text_encoder1, text_encoder2, vae, unet, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint(
-    #    "sdxl_base_v1-0", args.model_to_load, "cpu"
-    # )
-
-    # Implement Load model from ckpt or safetensors
     text_encoder1, text_encoder2, vae, unet = None, None, None, None
+    device = "cpu"
+    if is_valid_url(args.model_to_load):
+        tmp_model_name = "download"
+        download_dst_dir = tempfile.mkdtemp()
+        model_path = download_model(args.model_to_load, download_dst_dir, output_widget)
+        #model_path = os.path.join(download_dst_dir,tmp_model_name)
+        if model_path == None:
+            with output_widget:
+                print("Loading from Checkpoint failed, the request could not be completed")
+            return text_encoder1, text_encoder2, vae, unet
+        else:
+            # Implement Load model from ckpt or safetensors
+            try:
+                text_encoder1, text_encoder2, vae, unet, _, _ = load_models_from_sdxl_checkpoint(
+                    "sdxl_base_v1-0", model_path, device
+                )
+                return text_encoder1, text_encoder2, vae, unet
+            except Exception as e:
+                print(f"Could not load SDXL from checkpoint due to: \n{e}")
+                return text_encoder1, text_encoder2, vae, unet
 
-    with output_widget:
-        print("Loading from Checkpoint not implemented, please implement based on your model needs.")
+            with output_widget:
+                print(f"Loading from Checkpoint from URL needs to be implemented - using {model_path}")
+    else:
+        # Implement Load model from ckpt or safetensors
+        try:
+            text_encoder1, text_encoder2, vae, unet, _, _ = load_models_from_sdxl_checkpoint(
+                "sdxl_base_v1-0", args.model_to_load, device
+            )
+            return text_encoder1, text_encoder2, vae, unet
+        except Exception as e:
+            print(f"Could not load SDXL from checkpoint due to: \n{e}")
+            return text_encoder1, text_encoder2, vae, unet
+
+        with output_widget:
+            print("Loading from Checkpoint needs to be implemented.")
 
     return text_encoder1, text_encoder2, vae, unet
 
@@ -125,16 +283,14 @@ def convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, sav
 
 def save_sdxl_as_checkpoint(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget):
     """Saves the SDXL model components as a checkpoint file (placeholder)."""
-    # logit_scale = None
-    # ckpt_info = None
-
-    # key_count = sdxl_model_util.save_stable_diffusion_checkpoint(
-    #    args.model_to_save, text_encoder1, text_encoder2, unet, args.epoch, args.global_step, ckpt_info, vae, logit_scale, save_dtype
-    # )
+    logit_scale = None
+    ckpt_info = None
 
+    key_count = save_stable_diffusion_checkpoint(
+        args.model_to_save, text_encoder1, text_encoder2, unet, args.epoch, args.global_step, ckpt_info, vae, logit_scale, save_dtype
+        )
     with output_widget:
-        print("Saving as Checkpoint not implemented, please implement based on your model needs.")
-        # print(f"Model saved. Total converted state_dict keys: {key_count}")
+        print(f"Model saved. Total converted state_dict keys: {key_count}")
 
 def save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget):
     """Saves the SDXL model as a Diffusers model."""
@@ -170,7 +326,6 @@ def convert_model(model_to_load, save_precision_as, epoch, global_step, referenc
             self.output_path = output_path #Using output_path even if hardcoded
             self.fp16 = fp16
 
-    # Create a temporary directory for output
     with tempfile.TemporaryDirectory() as tmpdirname:
         args = Args(model_to_load, save_precision_as, epoch, global_step, reference_model, tmpdirname, fp16)
         args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")
@@ -246,8 +401,6 @@ def main(model_to_load, save_precision_as, epoch, global_step, reference_model,
   """Main function orchestrating the entire process."""
   output = gr.Markdown()
 
-  # Hardcode output_path
-  #output_path = "./converted_model" ##This is incorrect! This will save to current working directory, which isnt ideal.
   # Create tempdir, will only be there for the function
   with tempfile.TemporaryDirectory() as output_path:
     conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output)
@@ -263,6 +416,13 @@ with gr.Blocks() as demo:
     gr.Markdown(f"""
         ## **⚠️ IMPORTANT WARNINGS ⚠️**
         This App is Coded by an LLM partially, and for more information please go here: [Ktiseos Nyx](https://github.com/Ktiseos-Nyx/Sdxl-to-diffusers). The colab edition of this may indeed break AUP. This space is running on CPU and in theory SHOULD work, but may be slow. Earth and Dusk/ Ktiseos Nyx does not have the enterprise budget for ZERO GPU or any gpu sadly! Thank you to the community, John6666 especially for coming to aid when gemini would NOT fix the requirements. Support Ktiseos Nyx & Myself on Ko-fi: [![Ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/Z8Z8L4EO)
+
+        **Understanding the 'Model to Load' Input:**
+
+        This field can accept any of the following:
+        *   A Hugging Face model identifier (e.g., `stabilityai/stable-diffusion-xl-base-1.0`).
+        *   A direct URL to a .ckpt or .safetensors model file.
+        *   **Important:** Huggingface direct links need to end as /resolve/main/ and the name of the model after.
     """)
 
     model_to_load = gr.Textbox(label="Model to Load (Checkpoint or Diffusers)", placeholder="Path to model")
@@ -277,10 +437,9 @@ with gr.Blocks() as demo:
 
     reference_model = gr.Textbox(label="Reference Diffusers Model",
                                  placeholder="e.g., stabilityai/stable-diffusion-xl-base-1.0")
-    #output_path = gr.Textbox(label="Output Path", value="./converted_model") #Remove text box - using temp file approach
 
     gr.Markdown("## Hugging Face Hub Configuration")
-    hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Your Hugging Face write token")
+    hf_token = gr.Textbox(type="password", label="Hugging Face Token", placeholder="Your Hugging Face write token") #THIS IS NEEDED
     with gr.Row():
         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")