modified: app.py

app.py

@@ -39,6 +39,7 @@ def encode(init_image, torch_device, ae):
     init_image = init_image.to(torch_device)
     with torch.no_grad():
         init_image = ae.encode(init_image.to()).to(torch.bfloat16)
+    init_image = init_image.to(torch_device)
     return init_image
 
 
@@ -539,7 +540,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Flux")
     parser.add_argument("--name", type=str, default="flux-dev", choices=list(configs.keys()), help="Model name")
     parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
+    parser.add_argument("--offload", default=False, help="Offload model to CPU when not in use")
     parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
     parser.add_argument("--port", type=int, default=9090)
     args = parser.parse_args()