Create app.py

I HAVE NO CLUE WHAT IM DOING

app.py

@@ -0,0 +1,300 @@
+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 re
+import json
+import glob
+import gdown
+import requests
+import subprocess
+from urllib.parse import urlparse, unquote
+from pathlib import Path
+
+# ---------------------- DEPENDENCIES ----------------------
+
+def install_dependencies_gradio():
+    """Installs the necessary dependencies for the Gradio app.  Run this ONCE."""
+    try:
+        !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 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
+    else:
+        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'):
+        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
+    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 == "":
+        repo_id = user["name"] + "/" + model_name.strip()
+    else:
+        repo_id = orgs_name + "/" + model_name.strip()
+
+    try:
+        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")
+    except HfHubHTTPError as e:
+        print(f"Model repo '{repo_id}' exists, skipping create repo")
+
+    print(f"Model repo '{repo_id}' link: https://huggingface.co/{repo_id}\n")
+
+    return repo_id
+
+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"))
+
+# ---------------------- CONVERSION AND UPLOAD FUNCTIONS ----------------------
+
+def load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget):
+    """Loads the SDXL model from a checkpoint or Diffusers model."""
+    model_load_message = "checkpoint" if is_load_checkpoint else "Diffusers" + (" as fp16" if args.fp16 else "")
+    with output_widget:
+        print(f"Loading {model_load_message}: {args.model_to_load}")
+
+    if is_load_checkpoint:
+        loaded_model_data = load_from_sdxl_checkpoint(args, output_widget)
+    else:
+        loaded_model_data = load_sdxl_from_diffusers(args, load_dtype)
+
+    return loaded_model_data
+
+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
+
+    with output_widget:
+        print("Loading from Checkpoint not implemented, please implement based on your model needs.")
+
+    return text_encoder1, text_encoder2, vae, unet
+
+def load_sdxl_from_diffusers(args, load_dtype):
+    """Loads an SDXL model from a Diffusers model directory."""
+    pipeline = StableDiffusionXLPipeline.from_pretrained(
+        args.model_to_load, torch_dtype=load_dtype, tokenizer=None, tokenizer_2=None, scheduler=None
+    )
+    text_encoder1 = pipeline.text_encoder
+    text_encoder2 = pipeline.text_encoder_2
+    vae = pipeline.vae
+    unet = pipeline.unet
+
+    return text_encoder1, text_encoder2, vae, unet
+
+def convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, save_dtype, output_widget):
+    """Converts and saves the SDXL model as either a checkpoint or a Diffusers model."""
+    text_encoder1, text_encoder2, vae, unet = loaded_model_data
+    model_save_message = "checkpoint" + ("" if save_dtype is None else f" in {save_dtype}") if is_save_checkpoint else "Diffusers"
+
+    with output_widget:
+        print(f"Converting and saving as {model_save_message}: {args.model_to_save}")
+
+    if is_save_checkpoint:
+        save_sdxl_as_checkpoint(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget)
+    else:
+        save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget)
+
+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
+    # )
+
+    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}")
+
+def save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget):
+    """Saves the SDXL model as a Diffusers model."""
+    with output_widget:
+        reference_model_message = args.reference_model if args.reference_model is not None else 'default model'
+        print(f"Copying scheduler/tokenizer config from: {reference_model_message}")
+
+    # Save diffusers pipeline
+    pipeline = StableDiffusionXLPipeline(
+        vae=vae,
+        text_encoder=text_encoder1,
+        text_encoder_2=text_encoder2,
+        unet=unet,
+        scheduler=None,  # Replace None if there is a scheduler
+        tokenizer=None,  # Replace None if there is a tokenizer
+        tokenizer_2=None  # Replace None if there is a tokenizer_2
+    )
+
+    pipeline.save_pretrained(args.model_to_save)
+
+    with output_widget:
+        print(f"Model saved as {save_dtype}.")
+
+def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16, output_widget):
+    """Main conversion function."""
+    class Args:  # Defining Args locally within convert_model
+        def __init__(self, model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16):
+            self.model_to_load = model_to_load
+            self.save_precision_as = save_precision_as
+            self.epoch = epoch
+            self.global_step = global_step
+            self.reference_model = reference_model
+            self.output_path = output_path
+            self.fp16 = fp16
+
+    args = Args(model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16)
+    args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")
+
+    try:
+        load_dtype = torch.float16 if fp16 else None
+        save_dtype = get_save_dtype(save_precision_as)
+
+        is_load_checkpoint = determine_load_checkpoint(model_to_load)
+        is_save_checkpoint = not is_load_checkpoint  # reverse of load model
+
+        loaded_model_data = load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget)
+        convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, save_dtype, output_widget)
+
+        with output_widget:
+            return f"Conversion complete. Model saved to {args.model_to_save}"
+
+    except Exception as e:
+        with output_widget:
+            return f"Conversion failed: {e}"
+
+def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private, output_widget):
+    """Uploads a model to the Hugging Face Hub."""
+    try:
+        login(hf_token, add_to_git_credential=True)
+        api = HfApi()
+        user = api.whoami(hf_token)
+        model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)
+
+        # Determine upload parameters (adjust as needed)
+        path_in_repo = ""
+        trained_model = os.path.basename(model_path)
+
+        path_in_repo_local = path_in_repo if path_in_repo and not is_diffusers_model(model_path) else ""
+
+        notification = f"Uploading {trained_model} from {model_path} to https://huggingface.co/{model_repo}"
+        with output_widget:
+            print(notification)
+
+        if os.path.isdir(model_path):
+            if is_diffusers_model(model_path):
+                commit_message = f"Upload diffusers format: {trained_model}"
+                print("Detected diffusers model. Adjusting upload parameters.")
+            else:
+                commit_message = f"Upload checkpoint: {trained_model}"
+                print("Detected regular model. Adjusting upload parameters.")
+
+            api.upload_folder(
+                folder_path=model_path,
+                path_in_repo=path_in_repo_local,
+                repo_id=model_repo,
+                commit_message=commit_message,
+                ignore_patterns=".ipynb_checkpoints",
+            )
+        else:
+            commit_message = f"Upload file: {trained_model}"
+            api.upload_file(
+                path_or_fileobj=model_path,
+                path_in_repo=path_in_repo_local,
+                repo_id=model_repo,
+                commit_message=commit_message,
+            )
+        with output_widget:
+            return f"Model upload complete! Check it out at https://huggingface.co/{model_repo}/tree/main"
+
+    except Exception as e:
+        with output_widget:
+            return f"Upload failed: {e}"
+
+# ---------------------- GRADIO INTERFACE ----------------------
+
+def main(model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16, hf_token, orgs_name, model_name, make_private):
+  """Main function orchestrating the entire process."""
+  output = gr.Markdown()
+
+  conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16, output)
+
+  upload_output = upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private, output)
+
+  # Return a combined output
+  return f"{conversion_output}\n\n{upload_output}"
+
+with gr.Blocks() as demo:
+
+    # Add initial warnings (only once)
+    gr.Markdown("""
+        ## **⚠️ IMPORTANT WARNINGS ⚠️**
+        This app may violate Google Colab AUP.  Use at your own risk.  `xformers` may cause issues.
+    """)
+
+    model_to_load = gr.Textbox(label="Model to Load (Checkpoint or Diffusers)", placeholder="Path to model")
+    with gr.Row():
+        save_precision_as = gr.Dropdown(
+            choices=["fp16", "bf16", "float"], value="fp16", label="Save Precision As"
+        )
+        fp16 = gr.Checkbox(label="Load as fp16 (Diffusers only)")
+    with gr.Row():
+        epoch = gr.Number(value=0, label="Epoch to Write (Checkpoint)")
+        global_step = gr.Number(value=0, label="Global Step to Write (Checkpoint)")
+
+    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="/content/output")
+
+    gr.Markdown("## Hugging Face Hub Configuration")
+    hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Your Hugging Face write token")
+    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")
+    make_private = gr.Checkbox(label="Make Repository Private", value=False)
+
+    convert_button = gr.Button("Convert and Upload")
+    output = gr.Markdown()
+
+    convert_button.click(fn=main,
+                       inputs=[model_to_load, save_precision_as, epoch, global_step, reference_model,
+                               output_path, fp16, hf_token, orgs_name, model_name, make_private],
+                       outputs=output)
+
+demo.launch()