Upload app.py

app.py

@@ -638,7 +638,7 @@ class EmptyInitWrapper(torch.overrides.TorchFunctionMode):
         return func(*args, **kwargs)
 
 with EmptyInitWrapper():
-    model = Flux().to(dtype=torch.bfloat16, device="cuda")
+    model = Flux().to(dtype=torch.bfloat16, device=device)
     
     sd = load_file(f"{model_path}/consolidated_s6700.safetensors")
     sd = {k.replace("model.", ""): v for k, v in sd.items()}