Update app.py

app.py

@@ -3,15 +3,24 @@ import numpy as np
 import random
 import spaces
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, FluxPipeline
 
-from pruna import PrunaModel
+from pruna import PrunaModel, smash, SmashConfig
 
 dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "CPU"
 
-pipe = PrunaModel.from_hub(
-    "PrunaAI/FLUX.1-schnell-smashed"
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=dtype).to(device)
+smash_config = SmashConfig()
+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,
 )
 
 MAX_SEED = np.iinfo(np.int32).max