Update app.py

app.py

@@ -10,7 +10,7 @@ from pruna import PrunaModel, smash, SmashConfig
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "CPU"
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=dtype).to(device)
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=dtype).to("cpu")
 smash_config = SmashConfig()
 smash_config["factorizer"] = "qkv_diffusers"
 smash_config["cacher"] = "fora"
@@ -22,6 +22,8 @@ pipe = smash(
     model=pipe,
     smash_config=smash_config,
 )
+pipe.move_to_device(device)
+
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048