Update app.py

app.py

@@ -19,6 +19,7 @@ from safetensors.torch import load_file
 from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
 import gc
 from gradio_client import Client
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
@@ -27,8 +28,10 @@ os.environ["HF_HOME"] = cache_path
 
 
 torch.backends.cuda.matmul.allow_tf32 = True
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("ostris/Flex.1-alpha", subfolder="vae", torch_dtype=dtype).to(device)
 
-pipe = FluxPipeline.from_pretrained("AlekseyCalvin/HSTcolor_FlexSoonr", device_map=None, ignore_mismatched_sizes=True, torch_dtype=torch.bfloat16)
+pipe = DiffusionPipeline.from_pretrained("AlekseyCalvin/FlexAlpha_Scaled_Soonr", vae=taef1, ignore_mismatched_sizes=True, torch_dtype=torch.bfloat16)
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 model_id = ("zer0int/LongCLIP-GmP-ViT-L-14")
@@ -41,6 +44,8 @@ pipe.text_encoder = clip_model.text_model
 pipe.tokenizer_max_length = 248
 pipe.text_encoder.dtype = torch.bfloat16
 
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
@@ -100,6 +105,7 @@ def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height,
             width=width,
             height=height,
             generator=generator,
+            good_vae=good_vae,
             joint_attention_kwargs={"scale": lora_scale},
         ).images[0]
     return image