Update app.py

app.py

@@ -15,12 +15,12 @@ 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
-from accelerate import Accelerator  # Import accelerate
+from accelerate import Accelerator
 
 
 # ---------------------- DEPENDENCIES ----------------------
 def install_dependencies_gradio():
-    """Installs the necessary dependencies, including accelerate."""
+    """Installs the necessary dependencies."""
     try:
         subprocess.run(
             [
@@ -99,33 +99,39 @@ def download_model(model_path_or_url):
 
 
 def create_model_repo(api, user, orgs_name, model_name, make_private=False):
-    """Creates a Hugging Face model repository, handling missing inputs."""
+    """Creates a Hugging Face model repository, handling missing inputs and sanitizing the username."""
 
-    print("---- create_model_repo Called ----")  # Debug Print
-    print(f"  user: {user}")  # Debug Print
-    print(f"  orgs_name: {orgs_name}")  # Debug Print
-    print(f"  model_name: {model_name}")  # Debug Print
+    print("---- create_model_repo Called ----")
+    print(f"  user: {user}")
+    print(f"  orgs_name: {orgs_name}")
+    print(f"  model_name: {model_name}")
 
     if not model_name:
         model_name = f"converted-model-{datetime.now().strftime('%Y%m%d%H%M%S')}"
-        print(f"  Using default model_name: {model_name}")  # Debug Print
+        print(f"  Using default model_name: {model_name}")
 
     if orgs_name:
         repo_id = f"{orgs_name}/{model_name.strip()}"
     elif user:
-        repo_id = f"{user['name']}/{model_name.strip()}"
+        # --- Sanitize the username ---
+        sanitized_username = re.sub(r"[^a-zA-Z0-9._-]", "-", user['name'])  # Replace invalid chars with hyphens
+        print(f"  Original Username: {user['name']}") #Debugging
+        print(f"  Sanitized Username: {sanitized_username}") #Debugging
+        repo_id = f"{sanitized_username}/{model_name.strip()}"
     else:
-        raise ValueError("Must provide either an organization name or be logged in.")
+        raise ValueError(
+            "Must provide either an organization name or be logged in."
+        )
 
-    print(f"  repo_id: {repo_id}")  # Debug Print
+    print(f"  repo_id: {repo_id}")
 
     try:
         api.create_repo(repo_id=repo_id, repo_type="model", private=make_private)
         print(f"Model repo '{repo_id}' created.")
+        return repo_id
     except Exception as e:
         print(f"Error creating repo: {e}")
         raise
-    return repo_id
 
 def load_sdxl_checkpoint(checkpoint_path):
     """Loads checkpoint and extracts state dicts."""
@@ -207,22 +213,6 @@ def convert_and_save_sdxl_to_diffusers(checkpoint_path_or_url, output_path, refe
     pipeline.save_pretrained(output_path)
     print(f"Model saved as Diffusers format: {output_path}")
 
-def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private):
-    """Uploads a model to the Hugging Face Hub."""
-    print("---- upload_to_huggingface Called ----")  # 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
-    api = HfApi()
-    # --- CRUCIAL: Log in with the token FIRST ---
-    login(token=hf_token, add_to_git_credential=True)
-    user = api.whoami()  # Get the logged-in user *without* the token
-    print(f"  Logged-in user: {user}")  # Debug Print
-
-    model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)
-    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) ----------------------
 
 def main(
@@ -234,7 +224,7 @@ def main(
     model_name,
     make_private,
 ):
-    """Main function: SDXL checkpoint to Diffusers, with debugging prints."""
+    """Main function: SDXL checkpoint to Diffusers, always fp16."""
 
     print("---- Main Function Called ----")
     print(f"  model_to_load: {model_to_load}")
@@ -245,22 +235,63 @@ def main(
     print(f"  model_name: {model_name}")
     print(f"  make_private: {make_private}")
 
+    # --- Force Login at the Beginning of main() ---
     try:
-        convert_and_save_sdxl_to_diffusers(
-            model_to_load, output_path, reference_model
-        )
-        upload_to_huggingface(
-            output_path, hf_token, orgs_name, model_name, make_private
-        )
+        login(token=hf_token, add_to_git_credential=True)
+        api = HfApi()
+        user = api.whoami()  # Get logged-in user info
+        print(f"  Logged-in user: {user}")
+    except Exception as e:
+        error_message = f"Error during login: {e} Ensure a valid WRITE token is provided."
+        print(f"---- Main Function Error: {error_message} ----")
+        return error_message
+
+    # --- Strip Whitespace from Inputs ---
+    model_to_load = model_to_load.strip()
+    reference_model = reference_model.strip()
+    output_path = output_path.strip()
+    hf_token = hf_token.strip()  # Even though it's a password field, good practice
+    orgs_name = orgs_name.strip() if orgs_name else "" #Handle empty
+    model_name = model_name.strip() if model_name else "" #Handle empty
+
+    try:
+        convert_and_save_sdxl_to_diffusers(model_to_load, output_path, reference_model)
+
+        # --- Create Repo and Upload (Simplified) ---
+        if not model_name:
+            model_name = f"converted-model-{datetime.now().strftime('%Y%m%d%H%M%S')}"
+            print(f"Using default model_name: {model_name}")
+
+        if orgs_name:
+            repo_id = f"{orgs_name}/{model_name.strip()}"
+        elif user:
+            # Sanitize username here as well:
+            sanitized_username = re.sub(r"[^a-zA-Z0-9._-]", "-", user['name'])
+            repo_id = f"{sanitized_username}/{model_name.strip()}"
+
+        else:  # Should never happen because of login, but good practice
+            raise ValueError("Must provide either an organization name or be logged in.")
+        print(f"repo_id = {repo_id}")
+        try:
+            api.create_repo(repo_id=repo_id, repo_type="model", private=make_private)
+            print(f"Model repo '{repo_id}' created.")
+        except Exception as e:
+            print(f"Error in creating model repo: {e}")
+            raise
+
+        api.upload_folder(folder_path=output_path, repo_id=repo_id)
+        print(f"Model uploaded to: https://huggingface.co/{repo_id}")
+
         result = "Conversion and upload completed successfully!"
         print(f"---- Main Function Successful: {result} ----")
         return result
+
     except Exception as e:
         error_message = f"An error occurred: {e}"
         print(f"---- Main Function Error: {error_message} ----")
         return error_message
 
-# ---------------------- GRADIO INTERFACE (Corrected Button Placement) ----------------------
+# ---------------------- GRADIO INTERFACE ----------------------
 
 css = """
 #main-container {
@@ -286,15 +317,18 @@ with gr.Blocks(css=css) as demo:
     - Direct URLs to model files
     - Hugging Face model repositories (e.g., 'my-org/my-model' or 'my-org/my-model/file.safetensors')
 
+    ### How To Use 
+    - Insert URL or Repository Information into Field 1 SDXL Checkpoint (Path, URL, or HF Repo)
+    - Optional: Insert Reference Diffusers Model (Optional)
+    - Optional: Output Path (Diffusers Format)
+    - Insert your HF Token WRITE:  Get your HF token here: [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
+    - Organization Name (Optional)
+    - Create a Model Name for Output Purposes
+
     ### ℹ️ Important Notes:
     - This tool runs on **CPU**, conversion might be slower than on GPU.
     - For Hugging Face uploads, you need a **WRITE** token (not a read token).
-    - Get your HF token here: [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
-
-    ### 💾 Memory Usage:
     - This space is configured for **FP16** precision to reduce memory usage.
-    - Close other applications during conversion.
-    - For large models, ensure you have at least 16GB of RAM.
 
     ### 💻 Source Code:
     - [GitHub Repository](https://github.com/Ktiseos-Nyx/Gradio-SDXL-Diffusers)
@@ -307,11 +341,11 @@ 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)",
+                label="SDXL Checkpoint (Path, URL, or HF Repo)",  # Corrected Label
                 placeholder="Path, URL, or Hugging Face Repo ID (e.g., my-org/my-model or my-org/my-model/file.safetensors)",
             )
             reference_model = gr.Textbox(
-                label="Reference Diffusers Model (Optional)",
+                label="Reference Diffusers Model (Optional)",  # Corrected Label
                 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")
@@ -324,7 +358,6 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(variant="panel"):
             output = gr.Markdown(container=True)
 
-    # --- CORRECT BUTTON CLICK PLACEMENT ---
     convert_button.click(
         fn=main,
         inputs=[