Update app.py

app.py

@@ -13,16 +13,18 @@ import subprocess
 from urllib.parse import urlparse, unquote
 from pathlib import Path
 import tempfile
-from tqdm import tqdm
+#from tqdm import tqdm # Removed as not crucial and can break display in gradio.
 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 huggingface_hub import login, HfApi, hf_hub_download  # Import hf_hub_download
+from huggingface_hub.utils import validate_repo_id, HFValidationError
 from huggingface_hub.errors import HfHubHTTPError
 
+
 # ---------------------- DEPENDENCIES ----------------------
 def install_dependencies_gradio():
     """Installs the necessary dependencies."""
@@ -55,6 +57,49 @@ def create_model_repo(api, user, orgs_name, model_name, make_private=False):
     return repo_id
 
 # ---------------------- MODEL LOADING AND CONVERSION ----------------------
+def download_model(model_path_or_url):
+    """Downloads a model from a URL or Hugging Face Hub, handling various cases.
+    Args:
+        model_path_or_url: Can be a local path, a URL, a Hugging Face repo ID,
+            or a repo ID with a filename (e.g., "user/repo/file.safetensors").
+    Returns:
+        The local path to the downloaded (or existing) file.
+    """
+    try:
+        # Check if it's a valid Hugging Face repo ID (and potentially a file within)
+        try:
+            validate_repo_id(model_path_or_url)
+            # It's a valid repo ID; use hf_hub_download without a filename
+            local_path = hf_hub_download(repo_id=model_path_or_url)
+            return local_path
+        except HFValidationError:
+            # Not a simple repo ID.  Might be a repo ID with a filename, or a URL.
+            pass
+
+        if model_path_or_url.startswith("http://") or model_path_or_url.startswith("https://"):
+            # It's a URL, use hf_hub_download to handle it (it handles URLs gracefully).
+            local_path = hf_hub_download(repo_id=None, filename=None, repo_type=None, url=model_path_or_url)
+            return local_path
+        elif os.path.isfile(model_path_or_url): # Local File
+            return model_path_or_url
+        else: #HuggingFace Model
+             # Try splitting into repo ID and filename (for "user/repo/file.safetensors")
+            try:
+                parts = model_path_or_url.split("/", 1)  # Split only on the first /
+                if len(parts) == 2:
+                    repo_id, filename = parts
+                    validate_repo_id(repo_id) # Check the repo_id part.
+                    local_path = hf_hub_download(repo_id=repo_id, filename=filename)
+                    return local_path
+                else:
+                    raise ValueError("Invalid input")
+            except HFValidationError: #Still invalid
+                raise ValueError(f"Invalid model path or URL: {model_path_or_url}")
+
+
+    except Exception as e:
+        raise ValueError(f"Error downloading or accessing model: {e}")
+
 
 def load_sdxl_checkpoint(checkpoint_path):
     """Loads an SDXL checkpoint (.ckpt or .safetensors) and returns components."""
@@ -85,61 +130,57 @@ def load_sdxl_checkpoint(checkpoint_path):
 
 def build_diffusers_model(text_encoder1_state, text_encoder2_state, vae_state, unet_state, reference_model_path=None):
     """Builds the Diffusers pipeline components from the loaded state dicts."""
-    # --- Load configurations, create models (empty), load state dicts ---
+
+    # Default to SDXL base 1.0 if no reference model is provided
+    if not reference_model_path:
+        reference_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
 
     # 1. Text Encoders
-    if reference_model_path:
-        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")
-    else: #Default
-        config_text_encoder1 = CLIPTextConfig.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
-        config_text_encoder2 = CLIPTextConfig.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", 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")
 
     text_encoder1 = CLIPTextModel(config_text_encoder1)
     text_encoder2 = CLIPTextModel(config_text_encoder2)
     text_encoder1.load_state_dict(text_encoder1_state)
     text_encoder2.load_state_dict(text_encoder2_state)
-    text_encoder1.to(torch.float16)  # Ensure fp16
-    text_encoder2.to(torch.float16)
+    text_encoder1.to(torch.float16).to("cpu")  # Ensure fp16 and CPU
+    text_encoder2.to(torch.float16).to("cpu")
 
     # 2. VAE
-    if reference_model_path:
-          vae = AutoencoderKL.from_pretrained(reference_model_path, subfolder="vae")
-    else:
-          vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="vae")
+    vae = AutoencoderKL.from_pretrained(reference_model_path, subfolder="vae")
     vae.load_state_dict(vae_state)
-    vae.to(torch.float16)
+    vae.to(torch.float16).to("cpu")
 
     # 3. UNet
-    if reference_model_path:
-      unet = UNet2DConditionModel.from_pretrained(reference_model_path, subfolder="unet")
-    else:
-      unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet")
-
+    unet = UNet2DConditionModel.from_pretrained(reference_model_path, subfolder="unet")
     unet.load_state_dict(unet_state)
-    unet.to(torch.float16)
+    unet.to(torch.float16).to("cpu")
 
     return text_encoder1, text_encoder2, vae, unet
 
 
 
-def convert_and_save_sdxl_to_diffusers(checkpoint_path, output_path, reference_model_path):
-    """Converts an SDXL checkpoint to Diffusers format and saves it."""
+def convert_and_save_sdxl_to_diffusers(checkpoint_path_or_url, output_path, reference_model_path):
+    """Converts an SDXL checkpoint to Diffusers format and saves it.
+    Args:
+        checkpoint_path_or_url:  The path/URL/repo ID of the checkpoint.
+    """
+
+    # Download the model if necessary (handles URLs, repo IDs, and local paths)
+    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, text_encoder2, vae, unet = build_diffusers_model(text_encoder1_state, text_encoder2_state, vae_state, unet_state, reference_model_path)
 
 
-    pipeline = StableDiffusionXLPipeline(
-        vae=vae,
-        text_encoder=text_encoder1,
-        text_encoder_2=text_encoder2,
-        unet=unet,
-        # You'll likely need to add tokenizer, scheduler, etc., here from the reference model
-        tokenizer = pipeline.tokenizer,
-        tokenizer_2 = pipeline.tokenizer_2,
-        scheduler = pipeline.scheduler
-    )
+    # Load tokenizer and scheduler from the reference model
+    pipeline = StableDiffusionXLPipeline.from_pretrained(reference_model_path,
+                                                         text_encoder=text_encoder1,
+                                                         text_encoder_2=text_encoder2,
+                                                         vae=vae,
+                                                         unet=unet,
+                                                         torch_dtype=torch.float16,)
+    pipeline.to("cpu")
     pipeline.save_pretrained(output_path)
     print(f"Model saved as Diffusers format: {output_path}")
 
@@ -150,22 +191,25 @@ def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_priv
     api = HfApi()
     user = api.whoami(hf_token)
     model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)
-    api.upload_folder(folder_path=model_path, repo_id=model_repo)  # Use upload_folder
+    api.upload_folder(folder_path=model_path, repo_id=model_repo)
     print(f"Model uploaded to: https://huggingface.co/{model_repo}")
 
 # ---------------------- GRADIO INTERFACE ----------------------
 def main(model_to_load, reference_model, output_path, hf_token, orgs_name, model_name, make_private):
     """Main function: SDXL checkpoint to Diffusers, always fp16."""
 
-    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)
+    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)
+        return "Conversion and upload completed successfully!"
+    except Exception as e:
+        return f"An error occurred: {e}" # Return the error message
 
-    return "Conversion and upload completed successfully!"
 
 with gr.Blocks() as demo:
-    model_to_load = gr.Textbox(label="SDXL Checkpoint to Load (.ckpt or .safetensors)", placeholder="Path to checkpoint")
+    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)")
     reference_model = gr.Textbox(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="/content/output")  # Clarified label
+    output_path = gr.Textbox(label="Output Path (Diffusers Format)", value="output")  # Default changed to "output"
     hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Your Hugging Face write token")
     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")