Update app.py

app.py

@@ -21,7 +21,7 @@ import shutil
 import tempfile
 from functools import partial
 from optimum.quanto import quantize, qfloat8, freeze
-from diffusers import FluxPipeline
+from flux_inference import FluxPipeline
 
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
@@ -91,7 +91,7 @@ print('Quantizing, please wait...')
 quantize(pipe.transformer, qfloat8)  # Consider changing to qint8 if you get invalid values
 freeze(pipe.transformer)
 print('Model quantized!')
-pipe.enable_model_cpu_offload()
+pipe.to(device)
 
 
 # Load 3D generation models