Upload 2 files

Used Windsurf to fix a bunch of stuff

app.py

@@ -14,6 +14,12 @@ 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
 
 # ---------------------- UTILITY FUNCTIONS ----------------------
 
@@ -94,6 +100,7 @@ def is_diffusers_model(model_path):
     return required_folders.issubset(set(os.listdir(model_path))) and os.path.isfile(os.path.join(model_path, "model_index.json"))
 
 # ---------------------- 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)
@@ -314,38 +321,300 @@ def save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_d
     with output_widget:
         print(f"Model saved as {save_dtype}.")
 
-def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, 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 #Using output_path even if hardcoded
-            self.fp16 = fp16
-
-    with tempfile.TemporaryDirectory() as tmpdirname:
-        args = Args(model_to_load, save_precision_as, epoch, global_step, reference_model, tmpdirname, fp16)
-        args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")
+def get_save_dtype(precision):
+    """
+    Convert precision string to torch dtype
+    """
+    if precision == "float32" or precision == "fp32":
+        return torch.float32
+    elif precision == "float16" or precision == "fp16":
+        return torch.float16
+    elif precision == "bfloat16" or precision == "bf16":
+        return torch.bfloat16
+    else:
+        raise ValueError(f"Unsupported precision: {precision}")
 
-        try:
-            load_dtype = torch.float16 if fp16 else None
-            save_dtype = get_save_dtype(save_precision_as)
+def get_file_size(file_path):
+    """Get file size in GB."""
+    try:
+        size_bytes = Path(file_path).stat().st_size
+        return size_bytes / (1024 * 1024 * 1024)  # Convert to GB
+    except:
+        return None
 
-            is_load_checkpoint = determine_load_checkpoint(model_to_load)
-            is_save_checkpoint = not is_load_checkpoint  # reverse of load model
+def get_available_memory():
+    """Get available system memory in GB."""
+    return psutil.virtual_memory().available / (1024 * 1024 * 1024)
 
-            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)
+def estimate_memory_requirements(model_path, precision):
+    """Estimate memory requirements for model conversion."""
+    try:
+        # Base memory requirement for SDXL
+        base_memory = 8  # GB
+        
+        # Get model size if local file
+        model_size = get_file_size(model_path) if not is_valid_url(model_path) else None
+        
+        # Adjust for precision
+        memory_multiplier = 1.0 if precision in ["float16", "fp16", "bfloat16", "bf16"] else 2.0
+        
+        # Calculate total required memory
+        required_memory = (base_memory + (model_size if model_size else 12)) * memory_multiplier
+        
+        return required_memory
+    except:
+        return 16  # Default safe estimate
 
-            with output_widget:
-                return f"Conversion complete. Model saved to {args.model_to_save}"
+def validate_model(model_path, precision):
+    """
+    Validate the model before conversion.
+    Returns (is_valid, message)
+    """
+    try:
+        # Check if it's a URL
+        if is_valid_url(model_path):
+            try:
+                response = requests.head(model_path)
+                if response.status_code != 200:
+                    return False, "❌ Invalid URL or model not accessible"
+                if 'content-length' in response.headers:
+                    size_gb = int(response.headers['content-length']) / (1024 * 1024 * 1024)
+                    if size_gb < 0.1:
+                        return False, "❌ File too small to be a valid model"
+            except:
+                return False, "❌ Error checking URL"
+        
+        # Check if it's a local file
+        elif not model_path.startswith("stabilityai/") and not Path(model_path).exists():
+            return False, "❌ Model file not found"
+        
+        # Check available memory
+        available_memory = get_available_memory()
+        required_memory = estimate_memory_requirements(model_path, precision)
+        
+        if available_memory < required_memory:
+            return False, f"❌ Insufficient memory. Need {math.ceil(required_memory)}GB, but only {math.ceil(available_memory)}GB available"
+        
+        # Memory warning
+        memory_message = ""
+        if available_memory < required_memory * 1.5:
+            memory_message = "⚠️ Memory is tight. Consider closing other applications."
+        
+        return True, f"✅ Model validated successfully. {memory_message}"
+    
+    except Exception as e:
+        return False, f"❌ Validation error: {str(e)}"
 
-        except Exception as e:
-            with output_widget:
-                return f"Conversion failed: {e}"
+def cleanup_temp_files(directory=None):
+    """Clean up temporary files after conversion."""
+    try:
+        if directory:
+            shutil.rmtree(directory, ignore_errors=True)
+        # Clean up other temp files
+        temp_pattern = "*.tmp"
+        for temp_file in Path(".").glob(temp_pattern):
+            temp_file.unlink()
+    except Exception as e:
+        print(f"Warning: Error during cleanup: {e}")
+
+def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output_widget):
+    """Convert the model between different formats."""
+    temp_dir = None
+    history = ConversionHistory()
+    
+    try:
+        print("Starting model conversion...")
+        update_progress(output_widget, "⏳ Initializing conversion process...", 0)
+        
+        # Get optimization suggestions
+        available_memory = get_available_memory()
+        auto_suggestions = get_auto_optimization_suggestions(model_to_load, save_precision_as, available_memory)
+        history_suggestions = history.get_optimization_suggestions(model_to_load)
+        
+        # Display suggestions
+        if auto_suggestions or history_suggestions:
+            print("\n🔍 Optimization Suggestions:")
+            for suggestion in auto_suggestions + history_suggestions:
+                print(suggestion)
+            print("\n")
+        
+        # Validate model
+        is_valid, message = validate_model(model_to_load, save_precision_as)
+        if not is_valid:
+            raise ValueError(message)
+        print(message)
+        
+        args = SimpleNamespace()
+        args.model_to_load = model_to_load
+        args.save_precision_as = save_precision_as
+        args.epoch = epoch
+        args.global_step = global_step
+        args.reference_model = reference_model
+        args.fp16 = fp16
+
+        update_progress(output_widget, "🔍 Validating input model...", 10)
+        args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")
+        
+        save_dtype = get_save_dtype(save_precision_as)
+        
+        # Create temporary directory for processing
+        temp_dir = tempfile.mkdtemp(prefix="sdxl_conversion_")
+        
+        update_progress(output_widget, "📥 Loading model components...", 30)
+        is_load_checkpoint = determine_load_checkpoint(args.model_to_load)
+        if is_load_checkpoint is None:
+            raise ValueError("Invalid model format or path")
+            
+        update_progress(output_widget, "🔄 Converting model...", 50)
+        loaded_model_data = load_sdxl_model(args, is_load_checkpoint, save_dtype, output_widget)
+        
+        update_progress(output_widget, "💾 Saving converted model...", 80)
+        is_save_checkpoint = args.model_to_save.endswith(get_supported_extensions())
+        result = convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, save_dtype, output_widget)
+        
+        update_progress(output_widget, "✅ Conversion completed!", 100)
+        print(f"Model conversion completed. Saved to: {args.model_to_save}")
+        
+        # Verify the converted model
+        is_valid, verify_message = verify_model_structure(args.model_to_save)
+        if not is_valid:
+            raise ValueError(verify_message)
+        print(verify_message)
+        
+        # Record successful conversion
+        history.add_entry(
+            model_to_load,
+            {
+                'precision': save_precision_as,
+                'fp16': fp16,
+                'epoch': epoch,
+                'global_step': global_step
+            },
+            True,
+            "Conversion completed successfully"
+        )
+        
+        cleanup_temp_files(temp_dir)
+        return result
+        
+    except Exception as e:
+        if temp_dir:
+            cleanup_temp_files(temp_dir)
+        
+        # Record failed conversion
+        history.add_entry(
+            model_to_load,
+            {
+                'precision': save_precision_as,
+                'fp16': fp16,
+                'epoch': epoch,
+                'global_step': global_step
+            },
+            False,
+            str(e)
+        )
+        
+        error_msg = f"❌ Error during model conversion: {str(e)}"
+        print(error_msg)
+        return error_msg
+
+def update_progress(output_widget, message, progress):
+    """Update the progress bar and message in the UI."""
+    progress_bar = "▓" * (progress // 5) + "░" * ((100 - progress) // 5)
+    print(f"{message}\n[{progress_bar}] {progress}%")
+
+class ConversionHistory:
+    def __init__(self, history_file="conversion_history.json"):
+        self.history_file = history_file
+        self.history = self._load_history()
+
+    def _load_history(self) -> List[Dict]:
+        try:
+            with open(self.history_file, 'r') as f:
+                return json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return []
+
+    def _save_history(self):
+        with open(self.history_file, 'w') as f:
+            json.dump(self.history, f, indent=2)
+
+    def add_entry(self, model_path: str, settings: Dict, success: bool, message: str):
+        entry = {
+            'timestamp': datetime.now().isoformat(),
+            'model_path': model_path,
+            'settings': settings,
+            'success': success,
+            'message': message
+        }
+        self.history.append(entry)
+        self._save_history()
+
+    def get_optimization_suggestions(self, model_path: str) -> List[str]:
+        """Analyze history and provide optimization suggestions."""
+        suggestions = []
+        similar_conversions = [h for h in self.history if h['model_path'] == model_path]
+        
+        if similar_conversions:
+            success_rate = sum(1 for h in similar_conversions if h['success']) / len(similar_conversions)
+            if success_rate < 1.0:
+                failed_attempts = [h for h in similar_conversions if not h['success']]
+                if any('memory' in h['message'].lower() for h in failed_attempts):
+                    suggestions.append("⚠️ Previous attempts had memory issues. Consider using fp16 precision.")
+                if any('timeout' in h['message'].lower() for h in failed_attempts):
+                    suggestions.append("⚠️ Previous attempts timed out. Try breaking down the conversion process.")
+
+        return suggestions
+
+def verify_model_structure(model_path: str) -> tuple[bool, str]:
+    """Verify the structure of the converted model."""
+    try:
+        if model_path.endswith('.safetensors'):
+            # Verify safetensors structure
+            with safe_open(model_path, framework="pt") as f:
+                if not f.keys():
+                    return False, "❌ Invalid safetensors file: no tensors found"
+        
+        # Check for essential components
+        required_keys = ["model.diffusion_model", "first_stage_model"]
+        missing_keys = []
+        
+        # Load and check key components
+        state_dict = load_file(model_path)
+        for key in required_keys:
+            if not any(k.startswith(key) for k in state_dict.keys()):
+                missing_keys.append(key)
+        
+        if missing_keys:
+            return False, f"❌ Missing essential components: {', '.join(missing_keys)}"
+        
+        return True, "✅ Model structure verified successfully"
+    except Exception as e:
+        return False, f"❌ Model verification failed: {str(e)}"
+
+def get_auto_optimization_suggestions(model_path: str, precision: str, available_memory: float) -> List[str]:
+    """Generate automatic optimization suggestions based on model and system characteristics."""
+    suggestions = []
+    
+    # Memory-based suggestions
+    if available_memory < 16:
+        suggestions.append("💡 Limited memory detected. Consider these options:")
+        suggestions.append("   - Use fp16 precision to reduce memory usage")
+        suggestions.append("   - Close other applications before conversion")
+        suggestions.append("   - Use a machine with more RAM if available")
+    
+    # Precision-based suggestions
+    if precision == "float32" and available_memory < 32:
+        suggestions.append("💡 Consider using fp16 precision for better memory efficiency")
+    
+    # Model size-based suggestions
+    model_size = get_file_size(model_path) if not is_valid_url(model_path) else None
+    if model_size and model_size > 10:
+        suggestions.append("💡 Large model detected. Recommendations:")
+        suggestions.append("   - Ensure stable internet connection for URL downloads")
+        suggestions.append("   - Consider breaking down the conversion process")
+    
+    return suggestions
 
 def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private, output_widget):
     """Uploads a model to the Hugging Face Hub."""
@@ -428,53 +697,113 @@ def increment_filename(filename):
             return new_name
         counter += 1
 
-with gr.Blocks() as demo:
-
-    # Add initial warnings (only once)
-    gr.Markdown(f"""
-        ## **⚠️ IMPORTANT WARNINGS ⚠️**
-        This App is Coded by an LLM partially, and for more information please go here: [Ktiseos Nyx](https://github.com/Ktiseos-Nyx/Sdxl-to-diffusers). The colab edition of this may indeed break AUP. This space is running on CPU and in theory SHOULD work, but may be slow. Earth and Dusk/ Ktiseos Nyx does not have the enterprise budget for ZERO GPU or any gpu sadly! Thank you to the community, John6666 especially for coming to aid when gemini would NOT fix the requirements. 
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🎨 SDXL Model Converter
+    Convert SDXL models between different formats and precisions. Works on CPU!
+    
+    ### 📥 Input Sources Supported:
+    - Local model files (.safetensors, .ckpt, etc.)
+    - Direct URLs to model files
+    - Hugging Face model repositories (e.g., 'stabilityai/stable-diffusion-xl-base-1.0')
+    
+    ### ℹ️ Important Notes:
+    - This tool runs on CPU, though 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
+    
+    ### 💾 Memory Usage Tips:
+    - Use FP16 precision when possible to reduce memory usage
+    - Close other applications during conversion
+    - For large models, ensure you have at least 16GB of RAM
     """)
-    gr.HTML(
-    """
-    <a href='https://ko-fi.com/Z8Z8L4EO' target='_blank'>
-        <img height='36' style='border:0px;height:36px;' src='https://storage.ko-fi.com/cdn/kofi3.png?v=6' border='0' alt='Buy Me a Coffee at ko-fi.com' />
-    </a>
-    """
-)
-    gr.Markdown(f"""**Understanding the 'Model to Load' Input:**
-
-        This field can accept any of the following:
-        *   A Hugging Face model identifier (e.g., `stabilityai/stable-diffusion-xl-base-1.0`).
-        *   A direct URL to a .ckpt or .safetensors model file.
-        *   **Important:** Huggingface direct links need to end as /resolve/main/ and the name of the model after.""")
-    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)")
+        with gr.Column():
+            model_to_load = gr.Textbox(
+                label="Model Path/URL/HF Repo",
+                placeholder="Enter local path, URL, or Hugging Face model ID (e.g., stabilityai/stable-diffusion-xl-base-1.0)",
+                type="text"
+            )
+            
+            save_precision_as = gr.Dropdown(
+                choices=["float32", "float16", "bfloat16"],
+                value="float16",
+                label="Save Precision",
+                info="Choose model precision (float16 recommended for most cases)"
+            )
 
-    reference_model = gr.Textbox(label="Reference Diffusers Model",
-                                 placeholder="e.g., stabilityai/stable-diffusion-xl-base-1.0")
+            with gr.Row():
+                epoch = gr.Number(
+                    value=0,
+                    label="Epoch",
+                    precision=0,
+                    info="Optional: Set epoch number for the saved model"
+                )
+                global_step = gr.Number(
+                    value=0,
+                    label="Global Step",
+                    precision=0,
+                    info="Optional: Set training step for the saved model"
+                )
 
-    gr.Markdown("## Hugging Face Hub Configuration")
-    gr.Markdown("Your Access Token can be found here: [Token](https://huggingface.co/settings/tokens) **⚠️ IMPORTANT WARNINGS ⚠️** Do not share your token, and with reason you should refresh your token after use for security.")
-    hf_token = gr.Textbox(type="password", label="Hugging Face Token", placeholder="Your Hugging Face write token") #THIS IS NEEDED
-    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)
+            reference_model = gr.Textbox(
+                label="Reference Model (Optional)",
+                placeholder="Path to reference model for scheduler config",
+                info="Optional: Used to copy scheduler configuration"
+            )
+
+            fp16 = gr.Checkbox(
+                label="Load in FP16",
+                value=True,
+                info="Load model in half precision (recommended for CPU usage)"
+            )
+
+            # Hugging Face Upload Section
+            gr.Markdown("### Upload to Hugging Face (Optional)")
+            
+            hf_token = gr.Textbox(
+                label="Hugging Face Token",
+                placeholder="Enter your WRITE token from huggingface.co/settings/tokens",
+                type="password",
+                info=" Must be a WRITE token, not a read token!"
+            )
+
+            with gr.Row():
+                orgs_name = gr.Textbox(
+                    label="Organization Name",
+                    placeholder="Optional: Your organization name",
+                    info="Leave empty to use your personal account"
+                )
+                model_name = gr.Textbox(
+                    label="Model Name",
+                    placeholder="Name for your uploaded model",
+                    info="The name your model will have on Hugging Face"
+                )
 
-    convert_button = gr.Button("Convert and Upload")
-    output = gr.Markdown()
+            make_private = gr.Checkbox(
+                label="Make Private",
+                value=True,
+                info="Keep the uploaded model private on Hugging Face"
+            )
 
-    convert_button.click(fn=main,
-                       inputs=[model_to_load, save_precision_as, epoch, global_step, reference_model,
-                               fp16, hf_token, orgs_name, model_name, make_private],
-                       outputs=output)
+        with gr.Column():
+            output = gr.Markdown(label="Output")
+            convert_btn = gr.Button("Convert Model", variant="primary")
+            convert_btn.click(
+                fn=main,
+                inputs=[
+                    model_to_load,
+                    save_precision_as,
+                    epoch,
+                    global_step,
+                    reference_model,
+                    fp16,
+                    hf_token,
+                    orgs_name,
+                    model_name,
+                    make_private
+                ],
+                outputs=output
+            )
 
 demo.launch()

requirements.txt

@@ -1,12 +1,30 @@
-numpy==1.26.4
-torch>=2.1.0
-diffusers
-transformers
-einops
-open-clip-torch
-gradio
-safetensors
-accelerate
-requests
-tqdm
-gdown
+# Core dependencies
+numpy>=1.26.4
+torch>=2.0.0
+diffusers>=0.21.4
+transformers>=4.30.0
+einops>=0.7.0
+open-clip-torch>=2.23.0
+
+# UI and interface
+gradio>=3.50.2
+
+# Model handling
+safetensors>=0.3.1
+accelerate>=0.23.0
+
+# Utilities
+psutil>=5.9.0
+requests>=2.31.0
+tqdm>=4.65.0
+gdown>=4.7.1
+
+# Type checking and validation
+typing-extensions>=4.8.0
+pydantic>=2.0.0
+
+# File handling and compression
+fsspec>=2023.0.0
+filelock>=3.13.0
+
+# Note: This app is hosted on Hugging Face Spaces, so ensure compatibility with their environment.