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 diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
 
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
@@ -72,10 +72,12 @@ else:
     print("CUDA installation not found")
 
 
-device = torch.device('cuda')
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-base_model = "black-forest-labs/FLUX.1-dev"
-pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16, token=huggingface_token).to(device)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1, token=huggingface_token).to(device)
+torch.cuda.empty_cache()
 
 # Load and fuse LoRA BEFORE quantizing
 print('Loading and fusing lora, please wait...')
@@ -83,7 +85,6 @@ lora_path = hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.sa
 pipe.load_lora_weights(lora_path)
 pipe.fuse_lora(lora_scale=1.0)
 pipe.unload_lora_weights()
-pipe.to(device)
 pipe.enable_model_cpu_offload()