fix seed and export

StableDiffuser.py

@@ -63,7 +63,7 @@ class StableDiffuser(torch.nn.Module):
     def get_noise(self, batch_size, width, height, generator=None):
         param = list(self.parameters())[0]
         return torch.randn(
-            (batch_size, self.unet.in_channels, width // 8, height // 8),
+            (batch_size, self.unet.config.in_channels, width // 8, height // 8),
             generator=generator).type(param.dtype).to(param.device)
 
     def add_noise(self, latents, noise, step):

app.py

@@ -199,12 +199,13 @@ class Demo:
 
                         with gr.Column():
                             self.train_memory_options = gr.Markdown(interactive=False,
-                                value='Performance and VRAM usage optimizations, may not work on all devices.')
+                                value='Performance and VRAM usage optimizations, may not work on all devices:')
                             with gr.Row():
                                 self.train_use_adamw8bit_input = gr.Checkbox(label="8bit AdamW", value=True)
                                 self.train_use_xformers_input = gr.Checkbox(label="xformers", value=True)
                                 self.train_use_amp_input = gr.Checkbox(label="AMP", value=True)
-                                self.train_use_gradient_checkpointing_input = gr.Checkbox(label="Gradient checkpointing", value=True)
+                                self.train_use_gradient_checkpointing_input = gr.Checkbox(
+                                    label="Gradient checkpointing", value=False)
 
                     with gr.Column(scale=1):
 
@@ -248,9 +249,10 @@ class Demo:
                         )
 
                     with gr.Column(scale=1):
+                        self.export_status = gr.Button(
+                            value='', variant='primary', label='Status', interactive=False)
                         self.export_button = gr.Button(
-                            value="Export",
-                        )
+                            value="Export")
 
         self.infr_button.click(self.inference, inputs = [
             self.prompt_input_infr,
@@ -288,12 +290,12 @@ class Demo:
             self.save_path_input_export,
             self.save_half_export
         ],
-        outputs=[self.export_button]
+        outputs=[self.export_status]
         )
 
     def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr,
               use_adamw8bit=True, use_xformers=False, use_amp=False, use_gradient_checkpointing=False,
-              seed = -1,
+              seed=-1,
               pbar = gr.Progress(track_tqdm=True)):
 
         if self.training:
@@ -330,7 +332,7 @@ class Demo:
         try:
             self.training = True
             train(repo_id_or_path, img_size, prompt, modules, frozen, iterations, neg_guidance, lr, save_path,
-                use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing, seed=seed)
+                use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing, seed=int(seed))
         finally:
             self.training = False
 
@@ -355,8 +357,9 @@ class Demo:
         with finetuner:
             if save_half:
                 diffuser = diffuser.half()
-                diffuser.pipeline.to(torch.float16, torch_device=diffuser.device)
+                diffuser.pipeline.to('cpu', torch_dtype=torch.float16)
             diffuser.pipeline.save_pretrained(save_path)
+        return [gr.update(value=f'Done! Your model is at {save_path}.')]
 
 
     def inference(self, prompt, negative_prompt, seed, width, height, model_name, base_repo_id_or_path, pbar = gr.Progress(track_tqdm=True)):

train.py

@@ -52,7 +52,7 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
 
         if seed == -1:
             seed = random.randint(0, 2 ** 30)
-        set_seed(seed)
+        set_seed(int(seed))
 
         for i in pbar:
             with torch.no_grad():