Update app.py

app.py

@@ -2,77 +2,21 @@ import os
 import gradio as gr
 import torch
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL
-from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTextConfig, CLIPTokenizer
+from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTextConfig
 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 ----------------------
+# (download_model, create_model_repo, etc. - All unchanged, but included for completeness)
 
-
-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:
@@ -125,10 +69,21 @@ def download_model(model_path_or_url):
         raise ValueError(f"Error downloading or accessing model: {e}")
 
 
-
+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
 def load_sdxl_checkpoint(checkpoint_path):
-    """Loads checkpoint and extracts state dicts, handling Illustrious-xl."""
-
+    """Loads checkpoint and extracts state dicts."""
     if checkpoint_path.endswith(".safetensors"):
         state_dict = load_file(checkpoint_path, device="cpu")
     elif checkpoint_path.endswith(".ckpt"):
@@ -142,44 +97,34 @@ def load_sdxl_checkpoint(checkpoint_path):
     unet_state = OrderedDict()
 
     for key, value in state_dict.items():
-        if key.startswith("first_stage_model."):  # VAE
+        if key.startswith("first_stage_model."):
             vae_state[key.replace("first_stage_model.", "")] = value.to(torch.float16)
-        elif key.startswith("condition_model.model.text_encoder."):  # First Text Encoder
+        elif key.startswith("condition_model.model.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
+        elif key.startswith("condition_model.model.text_encoder_2."):
             text_encoder2_state[key.replace("condition_model.model.text_encoder_2.", "")] = value.to(torch.float16)
-        elif key.startswith("model.diffusion_model."):  # UNet
+        elif key.startswith("model.diffusion_model."):
             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
-):
+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_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
+    text_encoder2 = CLIPTextModelWithProjection(config_text_encoder2)  # Correct class
     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)
@@ -190,29 +135,16 @@ def build_diffusers_model(
     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."""
-
+def convert_and_save_sdxl_to_diffusers(checkpoint_path_or_url, output_path, reference_model_path):
+    """Converts and saves the 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_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,
+        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,
@@ -225,9 +157,6 @@ def convert_and_save_sdxl_to_diffusers(
     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)
@@ -237,8 +166,8 @@ def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_priv
     api.upload_folder(folder_path=model_path, repo_id=model_repo)
     print(f"Model uploaded to: https://huggingface.co/{model_repo}")
 
+# ---------------------- MAIN FUNCTION (with Debugging Prints) ----------------------
 
-# ---------------------- GRADIO INTERFACE ----------------------
 def main(
     model_to_load,
     reference_model,
@@ -248,7 +177,16 @@ def main(
     model_name,
     make_private,
 ):
-    """Main function: SDXL checkpoint to Diffusers, always fp16."""
+    """Main function: SDXL checkpoint to Diffusers, with debugging prints."""
+
+    print("---- Main Function Called ----")  # Debug Print
+    print(f"  model_to_load: {model_to_load}")  # Debug Print
+    print(f"  reference_model: {reference_model}")  # Debug Print
+    print(f"  output_path: {output_path}")  # Debug Print
+    print(f"  hf_token: {hf_token}")  # Debug Print
+    print(f"  orgs_name: {orgs_name}")  # Debug Print
+    print(f"  model_name: {model_name}")  # Debug Print
+    print(f"  make_private: {make_private}")  # Debug Print
 
     try:
         convert_and_save_sdxl_to_diffusers(
@@ -257,10 +195,15 @@ def main(
         upload_to_huggingface(
             output_path, hf_token, orgs_name, model_name, make_private
         )
-        return "Conversion and upload completed successfully!"
+        result = "Conversion and upload completed successfully!"
+        print(f"---- Main Function Successful: {result} ----")  # Debug Print
+        return result
     except Exception as e:
-        return f"An error occurred: {e}"  # Return the error message
+        error_message = f"An error occurred: {e}"
+        print(f"---- Main Function Error: {error_message} ----")  # Debug Print
+        return error_message
 
+# ---------------------- GRADIO INTERFACE (Corrected Button Placement) ----------------------
 
 css = """
 #main-container {
@@ -271,7 +214,7 @@ css = """
     color: #333;
 }
 #convert-button {
-    margin-top: 1em; /* Adds some space above the button */
+    margin-top: 1em;
 }
 """
 
@@ -306,7 +249,6 @@ with gr.Blocks(css=css) as demo:
 
     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)",
@@ -315,25 +257,17 @@ with gr.Blocks(css=css) as demo:
                 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"
-            )
+            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"
+        with gr.Column(variant="panel"):
             output = gr.Markdown(container=False)
 
+    # --- CORRECT BUTTON CLICK PLACEMENT ---
     convert_button.click(
         fn=main,
         inputs=[