Remove Python cache files and update .gitignore

- Delete all __pycache__/*.pyc files from src/ and src/ui/
- Add .huggingface/ to .gitignore to prevent committing cached tokens
- Update .gitignore with new patterns

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>

.gitignore

@@ -0,0 +1,36 @@
+# Python cache
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual environment
+venv/
+.venv/
+env/
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Environment variables (contains sensitive data)
+.env
+.env.local
+.env.*.local
+
+# Model cache (large files)
+.cache/
+
+# HuggingFace cache (contains tokens and configs)
+.huggingface/
+*.safetensors
+*.bin
+
+# Logs
+*.log
+logs/
+
+# OS files
+.DS_Store
+Thumbs.db

README.md

@@ -12,4 +12,37 @@ license: mit
 short_description: Merge SDXL checkpoints & LoRA and export with quantization
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SDXL Model Merger
+
+A Gradio-based web application for merging, generating with, and exporting Stable Diffusion XL (SDXL) checkpoints.
+
+## Features
+
+- **Load pipelines** from HuggingFace URLs with optional VAE and multiple LoRAs
+- **Generate images** with seamless tiling support for panoramic/360° outputs
+- **Export merged models** with quantization options (int8, int4, float8)
+
+## Usage on HuggingFace Spaces
+
+This app is optimized for both local and Space deployments:
+
+```bash
+# Local deployment
+python app.py
+
+# Space deployment with CPU fallback
+export DEPLOYMENT_ENV=spaces
+python app.py
+```
+
+For best results:
+- Use **GPU** (NVIDIA) for fast generation - ~8GB VRAM recommended
+- CPU mode is available but will be slower and use more RAM (~16GB+)
+
+## Requirements
+
+- Python 3.10+
+- PyTorch 2.0+
+- 4GB+ VRAM (GPU) or 16GB+ RAM (CPU)
+- ~2GB disk space for cached models
+

app.py

@@ -361,6 +361,8 @@ def create_app():
             return (
                 '<div class="status-warning">⏳ Loading started...</div>',
                 "Starting download...",
+                gr.update(interactive=False),
+                gr.update(visible=True)
             )
 
         def on_load_pipeline_complete(status_msg, progress_text):
@@ -369,19 +371,22 @@ def create_app():
                 return (
                     '<div class="status-success">✅ Pipeline loaded successfully!</div>',
                     progress_text,
-                    gr.update(interactive=True)
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
                 )
             elif "⚠️" in status_msg or "cancelled" in status_msg.lower():
                 return (
                     '<div class="status-warning">⚠️ Download cancelled</div>',
                     progress_text,
-                    gr.update(interactive=True)
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
                 )
             else:
                 return (
                     f'<div class="status-error">{status_msg}</div>',
                     progress_text,
-                    gr.update(interactive=True)
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
                 )
 
         # Cancel button for pipeline loading
@@ -390,24 +395,20 @@ def create_app():
         load_btn.click(
             fn=on_load_pipeline_start,
             inputs=[],
-            outputs=[load_status, load_progress],
-        ).then(
-            fn=lambda: (gr.update(visible=True), gr.update(interactive=False)),
-            inputs=[],
-            outputs=[cancel_load_btn, load_btn],
+            outputs=[load_status, load_progress, load_btn, cancel_load_btn],
         ).then(
             fn=load_pipeline,
             inputs=[checkpoint_url, vae_url, lora_urls, lora_strengths],
             outputs=[load_status, load_progress],
             show_progress="full",
         ).then(
-            fn=lambda: (gr.update(visible=False), gr.update(interactive=True)),
+            fn=lambda: (gr.update(interactive=True), gr.update(visible=False)),
             inputs=[],
-            outputs=[cancel_load_btn, load_btn],
+            outputs=[cancel_load_btn, cancel_load_btn],  # Just to hide it
         ).then(
             fn=on_load_pipeline_complete,
             inputs=[load_status, load_progress],
-            outputs=[load_status, load_progress, load_btn],
+            outputs=[load_status, load_progress, load_btn, cancel_load_btn],
         ).then(
             fn=lambda: (
                 gr.update(choices=["(None found)"] + get_cached_checkpoints()),

requirements.txt

@@ -19,6 +19,10 @@ gradio>=6.9.0
 # Download utilities
 tqdm>=4.65.0
 requests>=2.31.0
+huggingface-hub>=0.23.0
+
+# System monitoring (memory detection)
+psutil>=5.9.0
 
 # Optional: quantization support
 optimum-quanto>=0.2.0

src/config.py

@@ -1,6 +1,7 @@
 """Configuration constants and global settings for SDXL Model Merger."""
 
 import os
+import sys
 from pathlib import Path
 
 # ──────────────────────────────────────────────
@@ -11,11 +12,37 @@ CACHE_DIR = SCRIPT_DIR / ".cache"
 CACHE_DIR.mkdir(exist_ok=True)
 
 # ──────────────────────────────────────────────
-# Default URLs
+# Deployment Environment Detection
 # ──────────────────────────────────────────────
-DEFAULT_CHECKPOINT_URL = "https://civitai.com/api/download/models/354657?type=Model&format=SafeTensor&size=full&fp=fp16"
-DEFAULT_VAE_URL        = "https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl.vae.safetensors?download=true"
-DEFAULT_LORA_URLS      = "https://civitai.com/api/download/models/143197?type=Model&format=SafeTensor"
+DEPLOYMENT_ENV = os.environ.get("DEPLOYMENT_ENV", "local").lower()
+
+if DEPLOYMENT_ENV not in ("local", "spaces"):
+    print(f"⚠️ Unknown DEPLOYMENT_ENV '{DEPLOYMENT_ENV}', defaulting to 'local'")
+
+
+def is_running_on_spaces() -> bool:
+    """Check if running on HuggingFace Spaces."""
+    return DEPLOYMENT_ENV == "spaces"
+
+
+# ──────────────────────────────────────────────
+# Default URLs - Use HF models for Spaces compatibility
+# ──────────────────────────────────────────────
+DEFAULT_CHECKPOINT_URL = os.environ.get(
+    "DEFAULT_CHECKPOINT_URL",
+    "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors?download=true"
+)
+
+DEFAULT_VAE_URL = os.environ.get(
+    "DEFAULT_VAE_URL",
+    "https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl.vae.safetensors?download=true"
+)
+
+# Default LoRA - using HF instead of CivitAI
+DEFAULT_LORA_URLS = os.environ.get(
+    "DEFAULT_LORA_URLS",
+    "https://huggingface.co/nerijs/pixel-art-xl/resolve/main/pixel-art-xl.safetensors?download=true"
+)
 
 # ──────────────────────────────────────────────
 # PyTorch & Device Settings
@@ -26,27 +53,65 @@ import torch
 def get_device_info() -> tuple[str, str]:
     """
     Detect and return the optimal device for ML inference.
-    
+
     Returns:
         Tuple of (device_name, device_description)
     """
     if torch.cuda.is_available():
         device_name = "cuda"
-        device_desc = f"CUDA (GPU: {torch.cuda.get_device_name(0)})"
+        gpu_name = torch.cuda.get_device_name(0)
+        # Check available VRAM
+        try:
+            vram_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+            if vram_total < 8.0:
+                return device_name, f"CUDA (GPU: {gpu_name}, {vram_total:.1f}GB VRAM - low memory)"
+        except Exception:
+            pass
+        return device_name, f"CUDA (GPU: {gpu_name})"
+    
     elif torch.backends.mps.is_available():
         device_name = "mps"
-        device_desc = "Apple Silicon MPS"
-    else:
-        device_name = "cpu"
-        device_desc = "CPU (no GPU available)"
+        return device_name, "Apple Silicon MPS"
     
-    return device_name, device_desc
+    else:
+        # CPU fallback - check available RAM
+        try:
+            import psutil
+            ram_gb = psutil.virtual_memory().total / (1024**3)
+            if ram_gb < 16.0:
+                return "cpu", f"CPU (WARNING: {ram_gb:.1f}GB RAM - may be insufficient)"
+            return "cpu", f"CPU ({ram_gb:.1f}GB RAM)"
+        except Exception:
+            return "cpu", "CPU (no GPU available)"
 
 
 device, device_description = get_device_info()
 dtype = torch.float16
 
+# Check if we're on low-memory hardware and warn
+def check_memory_requirements() -> bool:
+    """Check if system meets minimum requirements. Returns True if OK."""
+    min_ram_gb = 8.0 if device == "cpu" else 4.0
+    
+    try:
+        import psutil
+        total_ram = psutil.virtual_memory().total / (1024**3)
+        
+        # On Spaces with CPU, RAM is limited - use float32 for safety
+        if is_running_on_spaces() and device == "cpu":
+            print(f"ℹ️ Spaces CPU mode detected: using float32 for stability")
+            return True
+        
+        if total_ram < min_ram_gb:
+            print(f"⚠️ Warning: Low memory ({total_ram:.1f}GB < {min_ram_gb}GB required)")
+            return False
+    except Exception:
+        pass
+    
+    return True
+
 print(f"🚀 Using device: {device_description}")
+check_memory_requirements()
 
 # ──────────────────────────────────────────────
 # Global State

src/downloader.py

@@ -14,6 +14,11 @@ def extract_model_id(url: str) -> str | None:
     return match.group(1) if match else None
 
 
+def is_huggingface_url(url: str) -> bool:
+    """Check if URL is a HuggingFace model download URL."""
+    return "huggingface.co" in url.lower()
+
+
 def get_safe_filename_from_url(
     url: str,
     default_name: str = "model.safetensors",
@@ -198,6 +203,8 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
 
     Checks for existing cached files before downloading. If a valid cache
     exists (file exists with matching expected size), skips re-download.
+    
+    Supports both HTTP(S) and HuggingFace Hub URLs.
 
     Args:
         url: File URL to download (http/https/file)
@@ -219,9 +226,9 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         if local_path.exists():
             import shutil
             output_path.parent.mkdir(parents=True, exist_ok=True)
-            
+
             print(f"  📁 Copying from cache: {local_path.name} → {output_path.name}")
-            
+
             # Copy the file to cache location
             shutil.copy2(str(local_path), str(output_path))
 

src/exporter.py

@@ -31,12 +31,18 @@ def export_merged_model(
         return None, "⚠️ Please load a pipeline first."
 
     try:
+        # Validate quantization type
+        valid_qtypes = ("none", "int8", "int4", "float8")
+        if qtype not in valid_qtypes:
+            return None, f"❌ Invalid quantization type: {qtype}. Must be one of: {valid_qtypes}"
+
         # Step 1: Unload LoRAs
         yield "💾 Exporting model...", "Unloading LoRAs..."
         if include_lora:
             try:
                 global_pipe.unload_lora_weights()
-            except Exception:
+            except Exception as e:
+                print(f"  ℹ️ Could not unload LoRAs: {e}")
                 pass
 
         merged_state_dict = {}

src/generator.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from .config import device, dtype, pipe as global_pipe
+from .config import device, dtype, pipe as global_pipe, is_running_on_spaces
 from .pipeline import enable_seamless_tiling
 
 
@@ -35,13 +35,27 @@ def generate_image(
     if not global_pipe:
         return None, "⚠️ Please load a pipeline first."
 
+    # For CPU mode, use float32 and warn about slow generation
+    effective_dtype = dtype
+    effective_device = device
+    
+    if is_running_on_spaces() and device == "cpu":
+        print("  ℹ️ CPU mode: using float32 for stability (generation will be slower)")
+        effective_dtype = torch.float32
+        # Update pipeline to use float32
+        global_pipe.unet.to(dtype=torch.float32)
+        global_pipe.text_encoder.to(dtype=torch.float32)
+        global_pipe.text_encoder_2.to(dtype=torch.float32)
+        if global_pipe.vae:
+            global_pipe.vae.to(dtype=torch.float32)
+
     # Enable seamless tiling on UNet & VAE decoder
     enable_seamless_tiling(global_pipe.unet, tile_x=tile_x, tile_y=tile_y)
     enable_seamless_tiling(global_pipe.vae.decoder, tile_x=tile_x, tile_y=tile_y)
 
     yield "🎨 Generating image...", f"Steps: 0/{steps} | CFG: {cfg}"
 
-    generator = torch.Generator(device=device).manual_seed(42)  # Fixed seed for reproducibility
+    generator = torch.Generator(device=effective_device).manual_seed(42)  # Fixed seed for reproducibility
     result = global_pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,

src/pipeline.py

@@ -9,8 +9,7 @@ from diffusers import (
     DPMSolverSDEScheduler,
 )
 
-from .config import device, dtype, pipe as global_pipe, CACHE_DIR, device_description
-from .downloader import download_file_with_progress, get_safe_filename_from_url, set_download_cancelled
+from .config import device, dtype, pipe as global_pipe, CACHE_DIR, device_description, is_running_on_spaces
 
 
 def _make_asymmetric_forward(module, pad_h: int, pad_w: int, tile_x: bool, tile_y: bool):
@@ -163,11 +162,23 @@ def load_pipeline(
         if progress:
             progress(0.3, desc="Loading text encoders...")
 
+        # For CPU/low-memory environments on Spaces, use device_map for better RAM management
+        load_kwargs = {
+            "torch_dtype": dtype,
+            "use_safetensors": True,
+            "safety_checker": None,
+        }
+
+        if is_running_on_spaces() and device == "cpu":
+            print("  ℹ️ CPU mode detected: enabling device_map='auto' for better RAM management")
+            load_kwargs["device_map"] = "auto"
+        else:
+            # For GPU, we'll move to device after loading
+            load_kwargs["variant"] = "fp16" if dtype == torch.float16 else None
+
         global_pipe = StableDiffusionXLPipeline.from_single_file(
             str(checkpoint_path),
-            torch_dtype=dtype,
-            use_safetensors=True,
-            safety_checker=None,
+            **load_kwargs,
         )
         print("  ✅ Text encoders loaded")
         
@@ -175,6 +186,12 @@ def load_pipeline(
             progress(0.5, desc="Loading UNet...")
 
         print("  ✅ UNet loaded")
+        
+        # Move to device (unless using device_map='auto' which handles this automatically)
+        if not is_running_on_spaces() or device != "cpu":
+            print(f"  ⚙️ Moving pipeline to device: {device_description}...")
+            global_pipe = global_pipe.to(device=device, dtype=dtype)
+        
         yield "⚙️ Pipeline loaded, setting up components...", f"Using device: {device_description}"
 
         # Load VAE into pipeline if provided