Update app.py

app.py

@@ -15,7 +15,8 @@ from pytorch_lightning import seed_everything
 from omegaconf import OmegaConf
 from einops import rearrange, repeat
 from tqdm import tqdm
-from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 import gradio as gr
 import shutil
 import tempfile
@@ -32,6 +33,8 @@ from src.utils.camera_util import (
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground, images_to_video
 
+from transformer_flux import FluxTransformer2DModel
+
 # Set up cache path
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
@@ -73,9 +76,29 @@ else:
 
 
 device = torch.device('cuda')
-
+# Initialize the base model
+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)
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=huggingface_token).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=huggingface_token).to(device)
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+    base_model,
+    vae=good_vae,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+    torch_dtype=dtype,
+    token=huggingface_token
+)
+
+MAX_SEED = 2**32 - 1
+
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 # Load and fuse LoRA BEFORE quantizing
 print('Loading and fusing lora, please wait...')
@@ -83,16 +106,7 @@ 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.transformer.to(device, dtype=torch.bfloat16)
-
-# Now quantize after LoRA is fused
-print('Quantizing, please wait...')
-# Try qint8 if qfloat8 produces invalid values
-quantize(pipe.transformer, qfloat8)  # Consider changing to qint8 if you get invalid values
-freeze(pipe.transformer)
-print('Model quantized!')
-pipe.to(device)
-
+pipe.transformer.to(device, dtype=torch.bfloat16)
 
 # Load 3D generation models
 config_path = 'configs/instant-mesh-large.yaml'