Spaces:
Paused
Paused
File size: 2,190 Bytes
38548f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from glob import glob
import os
import shutil
import gradio as gr
from infer.lib.train.process_ckpt import extract_small_model
def download_weight(exp_dir: str) -> str:
models = glob(f"{exp_dir}/G_*.pth")
if not models:
raise gr.Error("No model found")
latest_model = max(models, key=os.path.getctime)
print(f"Latest model: {latest_model}")
name = os.path.basename(exp_dir)
out = os.path.join(exp_dir, f"{name}.pth")
extract_small_model(
latest_model, out, "40k", True, "Model trained by ZeroGPU.", "v2"
)
return out
def download_expdir(exp_dir: str) -> str:
shutil.make_archive(exp_dir, "zip", exp_dir)
return f"{exp_dir}.zip"
def remove_expdir(exp_dir: str) -> str:
shutil.rmtree(exp_dir)
return ""
class ExportTab:
def __init__(self):
pass
def ui(self):
gr.Markdown("# Download Model or Experiment Directory")
gr.Markdown(
"You can download the latest model or the entire experiment directory here."
)
with gr.Row():
self.download_weight_btn = gr.Button(
value="Latest model (for inferencing)", variant="primary"
)
self.download_weight_output = gr.File(label="Prune latest model")
with gr.Row():
self.download_expdir_btn = gr.Button(
value="Download experiment directory", variant="primary"
)
self.download_expdir_output = gr.File(label="Archive experiment directory")
with gr.Row():
self.remove_expdir_btn = gr.Button(
value="REMOVE experiment directory", variant="stop"
)
def build(self, exp_dir: gr.Textbox):
self.download_weight_btn.click(
fn=download_weight,
inputs=[exp_dir],
outputs=[self.download_weight_output],
)
self.download_expdir_btn.click(
fn=download_expdir,
inputs=[exp_dir],
outputs=[self.download_expdir_output],
)
self.remove_expdir_btn.click(
fn=remove_expdir,
inputs=[exp_dir],
outputs=[exp_dir],
)
|