Update app.py

app.py

@@ -10,20 +10,9 @@ 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("cpu")
-smash_config = SmashConfig(device="cpu")
-smash_config["factorizer"] = "qkv_diffusers"
-smash_config["cacher"] = "fora"
-smash_config["fora_interval"] = 2
-smash_config["quantizer"] = "torchao"
-smash_config["torchao_quant_type"] = "fp8dq"
-smash_config["torchao_excluded_modules"] = "norm+embedding"
-pipe = smash(
-    model=pipe,
-    smash_config=smash_config,
+pipe = PrunaModel.from_hub(
+    "PrunaAI/Segmind-Vega-smashed"
 )
-pipe.move_to_device(device)
-
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048