feat: download infer pack

app/export.py

@@ -1,32 +1,80 @@
 from glob import glob
 import os
 import shutil
+import tempfile
 import gradio as gr
 from infer.lib.train.process_ckpt import extract_small_model
+from app.train import train_index
 
 
 def download_weight(exp_dir: str) -> str:
-    models = glob(f"{exp_dir}/G_*.pth")
-    if not models:
-        raise gr.Error("No model found")
+    checkpoints = glob(f"{exp_dir}/G_*.pth")
+    if not checkpoints:
+        raise gr.Error("No checkpoint found")
 
-    latest_model = max(models, key=os.path.getctime)
-    print(f"Latest model: {latest_model}")
+    latest_checkpoint = max(checkpoints, key=os.path.getctime)
+    print(f"Latest checkpoint: {latest_checkpoint}")
 
     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"
+        latest_checkpoint, out, "40k", True, "Model trained by ZeroGPU.", "v2"
     )
 
     return out
 
 
+def download_inference_pack(exp_dir: str) -> str:
+    net_g = download_weight(exp_dir)
+    index = glob(f"{exp_dir}/added_*.index")
+    if not index:
+        train_index(exp_dir)
+    index = glob(f"{exp_dir}/added_*.index")
+    if not index:
+        raise gr.Error("Index not found")
+
+    # make zip of those two files
+    tmp = os.path.join(exp_dir, "inference_pack")
+    if os.path.exists(tmp):
+        shutil.rmtree(tmp)
+    os.makedirs(tmp)
+    shutil.copy(net_g, tmp)
+    shutil.copy(index[0], tmp)
+    shutil.make_archive(tmp, "zip", tmp)
+    shutil.rmtree(tmp)
+
+    return f"{tmp}.zip"
+
+
 def download_expdir(exp_dir: str) -> str:
     shutil.make_archive(exp_dir, "zip", exp_dir)
     return f"{exp_dir}.zip"
 
 
+def remove_legacy_checkpoints(exp_dir: str):
+    checkpoints = glob(f"{exp_dir}/G_*.pth")
+    if not checkpoints:
+        raise gr.Error("No checkpoint found")
+
+    latest_checkpoint = max(checkpoints, key=os.path.getctime)
+    print(f"Latest checkpoint: {latest_checkpoint}")
+    for checkpoint in checkpoints:
+        if checkpoint != latest_checkpoint:
+            os.remove(checkpoint)
+            print(f"Removed: {checkpoint}")
+
+    checkpoints = glob(f"{exp_dir}/D_*.pth")
+    if not checkpoints:
+        raise gr.Error("No checkpoint found")
+
+    latest_checkpoint = max(checkpoints, key=os.path.getctime)
+    print(f"Latest checkpoint: {latest_checkpoint}")
+    for checkpoint in checkpoints:
+        if checkpoint != latest_checkpoint:
+            os.remove(checkpoint)
+            print(f"Removed: {checkpoint}")
+
+
 def remove_expdir(exp_dir: str) -> str:
     shutil.rmtree(exp_dir)
     return ""
@@ -48,12 +96,23 @@ class ExportTab:
             )
             self.download_weight_output = gr.File(label="Prune latest model")
 
+        with gr.Row():
+            self.download_inference_pack_btn = gr.Button(
+                value="Download inference pack (model + index)", variant="primary"
+            )
+            self.download_inference_pack_output = gr.File(label="Inference pack")
+
         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_legacy_checkpoints_btn = gr.Button(
+                value="Remove legacy checkpoints"
+            )
+
         with gr.Row():
             self.remove_expdir_btn = gr.Button(
                 value="REMOVE experiment directory", variant="stop"
@@ -66,12 +125,23 @@ class ExportTab:
             outputs=[self.download_weight_output],
         )
 
+        self.download_inference_pack_btn.click(
+            fn=download_inference_pack,
+            inputs=[exp_dir],
+            outputs=[self.download_inference_pack_output],
+        )
+
         self.download_expdir_btn.click(
             fn=download_expdir,
             inputs=[exp_dir],
             outputs=[self.download_expdir_output],
         )
 
+        self.remove_legacy_checkpoints_btn.click(
+            fn=remove_legacy_checkpoints,
+            inputs=[exp_dir],
+        )
+
         self.remove_expdir_btn.click(
             fn=remove_expdir,
             inputs=[exp_dir],