Update app.py

app.py

@@ -23,6 +23,32 @@ from typing import Dict, List, Optional
 from huggingface_hub import login, HfApi
 from types import SimpleNamespace
 
+# Remove unused imports
+# import os
+# import gradio as gr
+# import torch
+# 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 gdown
+# import requests
+# import subprocess
+# from urllib.parse import urlparse, unquote
+# from pathlib import Path
+# import tempfile
+# from tqdm import tqdm
+# import psutil
+# import math
+# import shutil
+# import hashlib
+# from datetime import datetime
+# from typing import Dict, List, Optional
+# from huggingface_hub import login, HfApi
+# from types import SimpleNamespace
+
 # ---------------------- UTILITY FUNCTIONS ----------------------
 
 def is_valid_url(url):
@@ -30,26 +56,34 @@ def is_valid_url(url):
     try:
         result = urlparse(url)
         return all([result.scheme, result.netloc])
-    except:
+    except Exception as e:
+        print(f"Error checking URL validity: {e}")
         return False
 
 def get_filename(url):
-    response = requests.get(url, stream=True)
-    response.raise_for_status()
+    """Extracts the filename from a URL."""
+    try:
+        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:
-        url_path = urlparse(url).path
-        filename = unquote(os.path.basename(url_path))
+        if 'content-disposition' in response.headers:
+            content_disposition = response.headers['content-disposition']
+            filename = re.findall('filename="?([^";]+)"?', content_disposition)[0]
+        else:
+            url_path = urlparse(url).path
+            filename = unquote(os.path.basename(url_path))
 
-    return filename
+        return filename
+    except Exception as e:
+        print(f"Error getting filename from URL: {e}")
+        return None
 
 def get_supported_extensions():
+    """Returns a tuple of supported model file extensions."""
     return tuple([".ckpt", ".safetensors", ".pt", ".pth"])
 
 def download_model(url, dst, output_widget):
+    """Downloads a model from a URL to the specified destination."""
     filename = get_filename(url)
     filepath = os.path.join(dst, filename)
     try:
@@ -60,32 +94,34 @@ def download_model(url, dst, output_widget):
                 if "/blob/" in url:
                     url = url.replace("/blob/", "/resolve/")
             subprocess.run(["aria2c","-x 16",url,"-d",dst,"-o",filename])
-        with output_widget:
-            return filepath
+        return filepath
     except Exception as e:
-       with output_widget:
-            return None
+        print(f"Error downloading model: {e}")
+        return None
 
 def determine_load_checkpoint(model_to_load):
     """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"}
-        if required_folders.issubset(set(os.listdir(model_to_load))) and os.path.isfile(os.path.join(model_to_load, "model_index.json")):
-            return False
+    try:
+        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"}
+            if required_folders.issubset(set(os.listdir(model_to_load))) and os.path.isfile(os.path.join(model_to_load, "model_index.json")):
+                return False
+    except Exception as e:
+        print(f"Error determining load checkpoint: {e}")
     return None  # handle this case as required
 
 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 == "":
-        repo_id = user["name"] + "/" + model_name.strip()
-    else:
-        repo_id = orgs_name + "/" + model_name.strip()
-
     try:
+        if orgs_name == "":
+            repo_id = user["name"] + "/" + model_name.strip()
+        else:
+            repo_id = orgs_name + "/" + model_name.strip()
+
         validate_repo_id(repo_id)
         api.create_repo(repo_id=repo_id, repo_type="model", private=make_private)
         print(f"Model repo '{repo_id}' didn't exist, creating repo")
@@ -98,46 +134,54 @@ def create_model_repo(api, user, orgs_name, model_name, make_private=False):
 
 def is_diffusers_model(model_path):
     """Checks if a given path is a valid Diffusers model directory."""
-    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"))
+    try:
+        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"))
+    except Exception as e:
+        print(f"Error checking if model is a Diffusers model: {e}")
+        return False
 
 # ---------------------- 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
+    try:
+        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
+    except Exception as e:
+        print(f"Error loading models from checkpoint: {e}")
+        return None
 
 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."""
@@ -665,7 +709,7 @@ def main(model_to_load, save_precision_as, epoch, global_step, reference_model,
 
   # 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, use_xformers, output)
+    conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, use_xformers, hf_token, orgs_name, model_name, make_private)
 
     upload_output = upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private)