Update app.py

app.py

@@ -14,6 +14,8 @@ import time
 from huggingface_hub import hf_hub_download
 from diffusers import FluxTransformer2DModel, FluxPipeline
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 import safetensors.torch
 from safetensors.torch import load_file
 from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
@@ -26,13 +28,20 @@ os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
+torch.set_float32_matmul_precision("medium")
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 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)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=bfloat16).to(device)
+good_vae = AutoencoderKL.from_pretrained("ostris/Flex.1-alpha", subfolder="vae", torch_dtype=bfloat16).to(device)
 
-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)
+dtype = torch.bfloat16
+base_model = "AlekseyCalvin/HSTcolor_FlexSoonr"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to("cuda")
+#pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
+torch.cuda.empty_cache()
 
 model_id = ("zer0int/LongCLIP-GmP-ViT-L-14")
 config = CLIPConfig.from_pretrained(model_id)
@@ -44,8 +53,7 @@ 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)
-
+pipe.vae = AutoencoderKL.from_pretrained("ostris/Flex.1-alpha", subfolder="vae", torch_dtype=dtype).to(device)
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f: