Update app.py

app.py

@@ -73,7 +73,9 @@ import torch.nn.functional as F
 from PIL import Image
 from torchvision import transforms
 from torchvision.transforms import v2
-from diffusers import HeunDiscreteScheduler
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from diffusers import FluxPipeline
 from pytorch_lightning import seed_everything
 import os
@@ -95,10 +97,14 @@ isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(d
 
 # model initialization and loading
 # flux
-flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device_0, dtype=torch.bfloat16)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device_0)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=access_token).to(device_0)
+# flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device_0, dtype=torch.bfloat16)
+flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, vae=taef1).to(device)
 flux_lora_ckpt_path = hf_hub_download(repo_id="LTT/xxx-ckpt", filename="rgb_normal_large.safetensors", repo_type="model")
 flux_pipe.load_lora_weights(flux_lora_ckpt_path)
 flux_pipe.to(device=device_0, dtype=torch.bfloat16)
+flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(flux_pipe)
 
 
 # lrm
@@ -210,7 +216,8 @@ def generate_multi_view_images(prompt, seed):
             width=resolution * 4,
             height=resolution * 2,
             output_type='np',
-            generator=generator
+            generator=generator,
+            good_vae=good_vae,
         ).images
     return images