ckpt download as well as torch bug fix

app.py

@@ -64,11 +64,12 @@ class Demo:
                     label="Learning Rate",
                     info='Learning rate used to train'
                 )
-                self.progress_bar = gr.Text(interactive=False, label="Training Progress")
 
                 self.train_button = gr.Button(
                     value="Train",
                 )
+
+                self.download = gr.Files()
                 
             with gr.Column(scale=2) as inference_column:
 
@@ -125,7 +126,7 @@ class Demo:
                         self.iterations_input,
                         self.lr_input
                     ],
-                    outputs=[self.train_button, self.infr_button, self.progress_bar]
+                    outputs=[self.train_button, self.infr_button, self.download]
                 )
 
     def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
@@ -201,6 +202,8 @@ class Demo:
             loss.backward()
             optimizer.step()
 
+        torch.save(finetuner.state_dict(), 'ft.ckpt')
+
         self.finetuner = finetuner.eval().half()
 
         self.diffuser = self.diffuser.eval().half()
@@ -209,7 +212,7 @@ class Demo:
 
         self.training = False
 
-        return [gr.update(interactive=True), gr.update(interactive=True), None]
+        return [gr.update(interactive=True), gr.update(interactive=True), 'ft.ckpt']
 
 
     def inference(self, prompt, seed, pbar = gr.Progress(track_tqdm=True)):

requirements.txt

@@ -1,6 +1,6 @@
 gradio
-torch
-torchvision
+torch==1.13.1 --index-url https://download.pytorch.org/whl/cu118
+torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu118
 diffusers
 transformers
 accelerate