Update app.py

app.py

@@ -93,7 +93,14 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
     if step_index == int(pipeline.num_timesteps * 0.5):
        # torch.set_float32_matmul_precision("medium")
         callback_kwargs["latents"] = callback_kwargs["latents"].to(torch.float64)
-        pipe.unet.to(torch.float64)
+        def change_dtype(module):
+            for child in module.children():
+                if len(list(child.children())) > 0:
+                    change_dtype(child)
+                for param in child.parameters():
+                    param.data = param.data.to(torch.float64)
+
+        change_dtype(pipeline.unet)
       #  pipe.guidance_scale=1.0
        # pipe.scheduler.set_timesteps(num_inference_steps*.70)
       #  print(f"-- setting step {pipeline.num_timesteps * 0.1} --")
@@ -104,7 +111,14 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         torch.backends.cuda.matmul.allow_tf32 = False
         torch.set_float32_matmul_precision("highest")
         callback_kwargs["latents"] = callback_kwargs["latents"].to(torch.bfloat16)
-        pipe.unet.to(torch.bfloat16)
+        def change_dtype(module):
+            for child in module.children():
+                if len(list(child.children())) > 0:
+                    change_dtype(child)
+                for param in child.parameters():
+                    param.data = param.data.to(torch.bfloat16)
+
+        change_dtype(pipeline.unet)
      #   pipe.vae = vae_a
      #   pipe.unet = unet_a
       #  torch.backends.cudnn.deterministic = False