FLUX.1-schnell-smashed

Running on Zero

davidberenstein1957 commited on 6 days ago

Commit

a8b3bdd

verified ·

1 Parent(s): 724679b

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -12,7 +12,8 @@ device = "cuda" if torch.cuda.is_available() else "CPU"
 pipe = PrunaModel.from_hub(
     "PrunaAI/Segmind-Vega-smashed"
-).move_to_device("cuda")
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048

 pipe = PrunaModel.from_hub(
     "PrunaAI/Segmind-Vega-smashed"
+)
+pipe.move_to_device(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048