lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-04-24 16:29:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

add a way to save models

Browse Source
This commit is contained in:
layerdiffusion
2024-08-19 06:30:49 -07:00
parent 4e8ba14dd0
commit 96f264ec6a
3 changed files with 68 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
modules/ui_checkpoint_merger.py
View File

@@ -1,4 +1,4 @@
import os
import gradio as gr
from modules import sd_models, sd_vae, errors, extras, call_queue
@@ -29,6 +29,38 @@ def modelmerger(*args):
class UiCheckpointMerger:
def __init__(self):
with gr.Blocks(analytics_enabled=False) as modelmerger_interface:
with gr.Accordion(open=True, label='Save Current Checkpoint (including all quantization)'):
with gr.Row():
textbox_file_name_forge = gr.Textbox(label="Filename (will save in /models/Stable-diffusion)", value='my_model.safetensors')
btn_save_unet_forge = gr.Button('Save UNet')
btn_save_ckpt_forge = gr.Button('Save Checkpoint')
with gr.Row():
result_html = gr.HTML('Ready to save ... (Currently only support saving Flux models)')
def save_unet(filename):
from modules.paths import models_path
long_filename = os.path.join(models_path, 'Stable-diffusion', filename)
os.makedirs(os.path.dirname(long_filename), exist_ok=True)
from modules import shared, sd_models
sd_models.forge_model_reload()
p = shared.sd_model.save_unet(long_filename)
print(f'Saved UNet at: {p}')
return f'Saved UNet at: {p}'
def save_checkpoint(filename):
from modules.paths import models_path
long_filename = os.path.join(models_path, 'Stable-diffusion', filename)
os.makedirs(os.path.dirname(long_filename), exist_ok=True)
from modules import shared
sd_models.forge_model_reload()
p = shared.sd_model.save_checkpoint(long_filename)
print(f'Saved checkpoint at: {p}')
return f'Saved checkpoint at: {p}'
btn_save_unet_forge.click(save_unet, inputs=textbox_file_name_forge, outputs=result_html)
btn_save_ckpt_forge.click(save_checkpoint, inputs=textbox_file_name_forge, outputs=result_html)
with gr.Row(equal_height=False):
with gr.Column(variant='compact'):
self.interp_description = gr.HTML(value=update_interp_description("Weighted sum"), elem_id="modelmerger_interp_description")