Update app.py

app.py

@@ -16,6 +16,7 @@ from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_
 from huggingface_hub import HfFileSystem, ModelCard
 from safetensors.torch import load_file
 from huggingface_hub import login
+from pathlib import Path
 
 hf_token = os.environ.get("HF_TOKEN")
 login(token=hf_token)
@@ -47,7 +48,10 @@ pipe.tokenizer = clip_processor.tokenizer
 pipe.text_encoder = clip_model.text_model
 pipe.tokenizer_max_length = 248
 pipe.text_encoder.dtype = torch.bfloat16
-    
+
+pipe.transformer.enable_layerwise_casting(
+            storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
+        )
 MAX_SEED = 2**32-1
 
 class calculateDuration: