pearsonkyle/SDXL-Model-Merger
2
1#!/usr/bin/env python32"""3SDXL Model Merger - Modernized with modular architecture and improved UI/UX.4 5This application allows you to:6- Load SDXL checkpoints with optional VAE and multiple LoRAs7- Generate images with seamless tiling support8- Export merged models with quantization options9 10Author: Qwen Code Assistant11"""12 13try:14 import spaces # noqa: F401 ā must be imported before torch/CUDA packages15except ImportError:16 pass17 18import gradio as gr19 20 21def create_app():22 """Create and configure the Gradio app."""23 24 header_css = """25 .header-gradient {26 background: linear-gradient(135deg, #10b981 0%, #7c3aed 100%);27 -webkit-background-clip: text;28 -webkit-text-fill-color: transparent;29 background-clip: text;30 }31 32 .feature-card {33 border-radius: 12px;34 padding: 20px;35 margin-bottom: 16px;36 box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);37 transition: transform 0.2s ease;38 }39 40 .feature-card:hover {41 transform: translateY(-2px);42 box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1);43 }44 45 .gradio-container .label {46 font-weight: 600;47 color: #374151;48 margin-bottom: 8px;49 }50 51 .status-success { color: #059669 !important; font-weight: 600; }52 .status-error { color: #dc2626 !important; font-weight: 600; }53 .status-warning { color: #d97706 !important; font-weight: 600; }54 55 .gradio-container .btn {56 border-radius: 8px;57 padding: 12px 24px;58 font-weight: 600;59 }60 61 .gradio-container textarea,62 .gradio-container input[type="number"],63 .gradio-container input[type="text"] {64 border-radius: 8px;65 border-color: #d1d5db;66 }67 68 .gradio-container textarea:focus,69 .gradio-container input:focus {70 outline: none;71 border-color: #6366f1;72 box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.1);73 }74 75 .gradio-container .tabitem {76 background: transparent;77 border-radius: 12px;78 }79 80 .progress-text {81 font-weight: 500;82 color: #6b7280 !important;83 }84 """85 86 from src.pipeline import load_pipeline87 from src.generator import generate_image88 from src.exporter import export_merged_model89 from src.config import get_cached_models, get_cached_checkpoints, get_cached_vaes, get_cached_loras90 91 with gr.Blocks(title="SDXL Model Merger") as demo:92 # Header section93 with gr.Column(elem_classes=["feature-card"]):94 gr.HTML("""95 <div style="text-align: center; margin-bottom: 24px;">96 <h1 style="font-size: 2.5em; margin: 0; line-height: 1.2;">97 <span class="header-gradient">SDXL Model Merger</span>98 </h1>99 <p style="color: #6b7280; font-size: 1.1em; max-width: 600px; margin: 16px auto;">100 Merge checkpoints, LoRAs, and VAEs - then bake LoRAs into a single exportable101 checkpoint with optional quantization.102 </p>103 </div>104 """)105 106 # Feature highlights107 with gr.Row():108 with gr.Column(scale=1):109 gr.HTML("""110 <div style="text-align: center; padding: 16px;">111 <div style="font-size: 2.5em; margin-bottom: 8px;">š</div>112 <strong>Fast Loading</strong>113 <p style="font-size: 0.85em; color: #6b7280; margin-top: 4px;">With progress tracking & cache</p>114 </div>115 """)116 with gr.Column(scale=1):117 gr.HTML("""118 <div style="text-align: center; padding: 16px;">119 <div style="font-size: 2.5em; margin-bottom: 8px;">šØ</div>120 <strong>Panorama Gen</strong>121 <p style="font-size: 0.85em; color: #6b7280; margin-top: 4px;">Seamless tiling support</p>122 </div>123 """)124 with gr.Column(scale=1):125 gr.HTML("""126 <div style="text-align: center; padding: 16px;">127 <div style="font-size: 2.5em; margin-bottom: 8px;">š¦</div>128 <strong>Export Ready</strong>129 <p style="font-size: 0.85em; color: #6b7280; margin-top: 4px;">Quantization & format options</p>130 </div>131 """)132 133 gr.Markdown("---")134 135 with gr.Tab("Load Pipeline"):136 gr.Markdown("### Load SDXL Pipeline with Checkpoint, VAE, and LoRAs")137 138 # Progress indicator for pipeline loading139 load_progress = gr.Textbox(140 label="Loading Progress",141 placeholder="Ready to start...",142 show_label=True,143 info="Real-time status of model downloads and pipeline setup"144 )145 146 with gr.Row():147 with gr.Column(scale=2):148 # Checkpoint URL with cached models dropdown149 checkpoint_url = gr.Textbox(150 label="Base Model (.safetensors) URL",151 value="https://civitai.com/api/download/models/354657?type=Model&format=SafeTensor&size=full&fp=fp16",152 placeholder="e.g., https://civitai.com/api/download/models/...",153 info="Download link for the base SDXL checkpoint"154 )155 156 # Dropdown of cached checkpoints157 cached_checkpoints = gr.Dropdown(158 choices=["(None found)"] + get_cached_checkpoints(),159 label="Cached Checkpoints",160 value="(None found)" if not get_cached_checkpoints() else None,161 info="Models already downloaded to .cache/"162 )163 164 # VAE URL165 vae_url = gr.Textbox(166 label="VAE (.safetensors) URL",167 value="https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl.vae.safetensors?download=true",168 placeholder="Leave blank to use model's built-in VAE",169 info="Optional custom VAE for improved quality"170 )171 172 # Dropdown of cached VAEs173 cached_vaes = gr.Dropdown(174 choices=["(None found)"] + get_cached_vaes(),175 label="Cached VAEs",176 value="(None found)" if not get_cached_vaes() else None,177 info="Select a VAE to load"178 )179 180 with gr.Column(scale=1):181 # LoRA URLs input182 lora_urls = gr.Textbox(183 label="LoRA URLs (one per line)",184 lines=5,185 value="https://civitai.com/api/download/models/143197?type=Model&format=SafeTensor",186 placeholder="https://civit.ai/...\nhttps://huggingface.co/...",187 info="Multiple LoRAs can be loaded and fused together"188 )189 190 # Dropdown of cached LoRAs191 cached_loras = gr.Dropdown(192 choices=["(None found)"] + get_cached_loras(),193 label="Cached LoRAs",194 value="(None found)" if not get_cached_loras() else None,195 info="Select a LoRA to add to the list below"196 )197 198 lora_strengths = gr.Textbox(199 label="LoRA Strengths",200 value="1.0",201 placeholder="e.g., 0.8,1.0,0.5",202 info="Comma-separated strength values for each LoRA"203 )204 205 with gr.Row():206 load_btn = gr.Button("š Load Pipeline", variant="primary", size="lg")207 208 # Detailed status display209 load_status = gr.HTML(210 label="Status",211 value='<div class="status-success">ā
Ready to load pipeline</div>',212 )213 214 with gr.Tab("Generate Image"):215 gr.Markdown("### Generate Panorama Images with Seamless Tiling")216 217 # Progress indicator for image generation218 gen_progress = gr.Textbox(219 label="Generation Progress",220 placeholder="Ready to generate...",221 show_label=True,222 info="Real-time status of image generation"223 )224 225 with gr.Row():226 with gr.Column(scale=1):227 prompt = gr.Textbox(228 label="Positive Prompt",229 value="Glowing mushrooms around pyramids amidst a cosmic backdrop, equirectangular, 360 panorama, cinematic",230 lines=4,231 placeholder="Describe the image you want to generate..."232 )233 234 cfg = gr.Slider(235 minimum=1.0, maximum=20.0, value=3.0, step=0.5,236 label="CFG Scale",237 info="Higher values make outputs match prompt more strictly"238 )239 240 height = gr.Number(241 value=1024, precision=0,242 label="Height (pixels)",243 info="Output image height"244 )245 246 with gr.Column(scale=1):247 negative_prompt = gr.Textbox(248 label="Negative Prompt",249 value="boring, text, signature, watermark, low quality, bad quality",250 lines=4,251 placeholder="Elements to avoid in generation..."252 )253 254 steps = gr.Slider(255 minimum=1, maximum=100, value=8, step=1,256 label="Inference Steps",257 info="More steps = better quality but slower"258 )259 260 width = gr.Number(261 value=2048, precision=0,262 label="Width (pixels)",263 info="Output image width"264 )265 266 with gr.Row():267 tile_x = gr.Checkbox(True, label="X-axis Seamless Tiling")268 tile_y = gr.Checkbox(False, label="Y-axis Seamless Tiling")269 270 seed = gr.Number(271 value=80484030936239,272 precision=0,273 label="Seed",274 info="Random seed for reproducible generation"275 )276 277 with gr.Row():278 gen_btn = gr.Button("āØ Generate Image", variant="secondary", size="lg")279 280 with gr.Row():281 image_output = gr.Image(282 label="Result",283 height=400,284 show_label=True285 )286 with gr.Column():287 gen_status = gr.HTML(288 label="Generation Status",289 value='<div class="status-success">ā
Ready to generate</div>',290 )291 292 gr.HTML("""293 <div style="margin-top: 16px; padding: 12px; background-color: #e5e7eb !important; border-radius: 8px;">294 <strong style="color: #1f2937 !important;">š” Tips:</strong>295 <ul style="margin: 8px 0; padding-left: 20px; font-size: 0.9em; color: #1f2937 !important;">296 <li>Use wide aspect ratios (e.g., 1024x2048) for panoramas</li>297 <li>Enable seamless tiling for texture-like outputs</li>298 <li>Lower CFG (3-5) for more creative results</li>299 </ul>300 </div>301 """)302 303 with gr.Tab("Export Model"):304 gr.Markdown("### Export Merged Checkpoint with Quantization Options")305 306 # Progress indicator for export307 export_progress = gr.Textbox(308 label="Export Progress",309 placeholder="Ready to export...",310 show_label=True,311 info="Real-time status of model export and quantization"312 )313 314 with gr.Row():315 include_lora = gr.Checkbox(316 True,317 label="Include Fused LoRAs",318 info="Bake the loaded LoRAs into the exported model"319 )320 321 quantize_toggle = gr.Checkbox(322 False,323 label="Apply Quantization",324 info="Reduce model size with quantization"325 )326 327 qtype_row = gr.Row(visible=True)328 with qtype_row:329 qtype_dropdown = gr.Dropdown(330 choices=["none", "int8", "int4", "float8"],331 value="int8",332 label="Quantization Method",333 info="Trade quality for smaller file size"334 )335 336 with gr.Row():337 format_dropdown = gr.Dropdown(338 choices=["safetensors", "bin"],339 value="safetensors",340 label="Export Format",341 info="safetensors is recommended for safety"342 )343 344 with gr.Row():345 export_btn = gr.Button("š¾ Save Merged Checkpoint", variant="primary", size="lg")346 347 with gr.Row():348 download_link = gr.File(349 label="Download Merged File",350 show_label=True,351 )352 353 with gr.Column():354 export_status = gr.HTML(355 label="Export Status",356 value='<div class="status-success">ā
Ready to export</div>',357 )358 359 gr.HTML("""360 <div style="margin-top: 16px; padding: 12px; background: #e0f2fe; border-radius: 8px;">361 <strong>ā¹ļø About Quantization:</strong>362 <p style="font-size: 0.9em; margin: 8px 0;">363 Reduces model size by lowering precision. Int8 is typically364 lossless for inference while cutting size in half.365 </p>366 </div>367 """)368 369 # Event handlers - all inside Blocks context370 371 def on_load_pipeline_start():372 """Called when pipeline loading starts."""373 return (374 '<div class="status-warning">ā³ Loading started...</div>',375 "Starting download...",376 gr.update(interactive=False)377 )378 379 def on_load_pipeline_complete(status_msg, progress_text):380 """Called when pipeline loading completes."""381 if "ā
" in status_msg:382 return (383 '<div class="status-success">ā
Pipeline loaded successfully!</div>',384 progress_text,385 gr.update(interactive=True)386 )387 elif "ā ļø" in status_msg or "cancelled" in status_msg.lower():388 return (389 '<div class="status-warning">ā ļø Download cancelled</div>',390 progress_text,391 gr.update(interactive=True)392 )393 else:394 return (395 f'<div class="status-error">{status_msg}</div>',396 progress_text,397 gr.update(interactive=True)398 )399 400 load_btn.click(401 fn=on_load_pipeline_start,402 inputs=[],403 outputs=[load_status, load_progress, load_btn],404 ).then(405 fn=load_pipeline,406 inputs=[checkpoint_url, vae_url, lora_urls, lora_strengths],407 outputs=[load_status, load_progress],408 show_progress="full",409 ).then(410 fn=on_load_pipeline_complete,411 inputs=[load_status, load_progress],412 outputs=[load_status, load_progress, load_btn],413 ).then(414 fn=lambda: (415 gr.update(choices=["(None found)"] + get_cached_checkpoints()),416 gr.update(choices=["(None found)"] + get_cached_vaes()),417 gr.update(choices=["(None found)"] + get_cached_loras()),418 ),419 inputs=[],420 outputs=[cached_checkpoints, cached_vaes, cached_loras],421 )422 423 def on_cached_checkpoint_change(cached_path):424 """Update URL when a cached checkpoint is selected."""425 if cached_path and cached_path != "(None found)":426 return gr.update(value=f"file://{cached_path}")427 return gr.update()428 429 cached_checkpoints.change(430 fn=lambda x: gr.update(value=f"file://{x}" if x and x != "(None found)" else ""),431 inputs=cached_checkpoints,432 outputs=checkpoint_url,433 )434 435 def on_cached_vae_change(cached_path):436 """Update VAE URL when a cached VAE is selected."""437 if cached_path and cached_path != "(None found)":438 return gr.update(value=f"file://{cached_path}")439 return gr.update()440 441 cached_vaes.change(442 fn=on_cached_vae_change,443 inputs=cached_vaes,444 outputs=vae_url,445 )446 447 def on_cached_lora_change(cached_path, current_urls):448 """Add cached LoRA to the list."""449 if cached_path and cached_path != "(None found)":450 urls_list = [u.strip() for u in current_urls.split("\n") if u.strip()]451 file_url = f"file://{cached_path}"452 if file_url not in urls_list:453 urls_list.append(file_url)454 return gr.update(value="\n".join(urls_list))455 return gr.update()456 457 cached_loras.change(458 fn=on_cached_lora_change,459 inputs=[cached_loras, lora_urls],460 outputs=lora_urls,461 )462 463 464 def on_generate_start():465 """Called when image generation starts."""466 return (467 '<div class="status-warning">ā³ Generating image...</div>',468 "Starting generation...",469 gr.update(interactive=False)470 )471 472 def on_generate_complete(status_msg, progress_text, image):473 """Called when image generation completes."""474 if image is None:475 return (476 f'<div class="status-error">{status_msg}</div>',477 "",478 gr.update(interactive=True),479 gr.update()480 )481 else:482 return (483 '<div class="status-success">ā
Generation complete!</div>',484 "Done",485 gr.update(interactive=True),486 gr.update(value=image)487 )488 489 gen_btn.click(490 fn=on_generate_start,491 inputs=[],492 outputs=[gen_status, gen_progress, gen_btn],493 ).then(494 fn=generate_image,495 inputs=[prompt, negative_prompt, cfg, steps, height, width, tile_x, tile_y, seed],496 outputs=[image_output, gen_progress],497 ).then(498 fn=lambda img, msg: on_generate_complete(msg, "Done", img),499 inputs=[image_output, gen_progress],500 outputs=[gen_status, gen_progress, gen_btn, image_output],501 )502 503 def on_export_start():504 """Called when export starts."""505 return (506 '<div class="status-warning">ā³ Export started...</div>',507 "Starting export...",508 gr.update(interactive=False)509 )510 511 def on_export_complete(status_msg, progress_text, file_path):512 """Called when export completes."""513 if file_path is None:514 return (515 f'<div class="status-error">{status_msg}</div>',516 "",517 gr.update(interactive=True),518 gr.update(value=None)519 )520 else:521 return (522 '<div class="status-success">ā
Export complete!</div>',523 "Exported successfully",524 gr.update(interactive=True),525 gr.update(value=file_path)526 )527 528 export_btn.click(529 fn=on_export_start,530 inputs=[],531 outputs=[export_status, export_progress, export_btn],532 ).then(533 fn=lambda inc, q, qt, fmt: export_merged_model(534 include_lora=inc,535 quantize=q and (qt != "none"),536 qtype=qt, # always pass the string value; exporter handles "none" correctly537 save_format=fmt,538 ),539 inputs=[include_lora, quantize_toggle, qtype_dropdown, format_dropdown],540 outputs=[download_link, export_progress],541 ).then(542 fn=lambda path, msg: on_export_complete(msg, "Exported", path),543 inputs=[download_link, export_progress],544 outputs=[export_status, export_progress, export_btn, download_link],545 )546 547 quantize_toggle.change(548 fn=lambda checked: gr.update(visible=checked),549 inputs=[quantize_toggle],550 outputs=qtype_row,551 )552 553 return demo554 555 556demo = create_app()557 558if __name__ == "__main__":559 demo.launch()560