Update app.py

app.py

@@ -13,7 +13,7 @@ import random
 import time
 from typing import Any, Dict, List, Optional, Union
 from huggingface_hub import hf_hub_download
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoPipelineForImage2Image, ConfigMixin, FluxTransformer2DModel
+from diffusers import DiffusionPipeline, AutoencoderTiny, ModelMixin, AutoPipelineForImage2Image, ConfigMixin, FluxTransformer2DModel
 import safetensors.torch
 from safetensors.torch import load_file
 from pipeline import FluxWithCFGPipeline
@@ -38,8 +38,12 @@ with open('loras.json', 'r') as f:
 
 dtype = torch.bfloat16
 
-pipe = FluxWithCFGPipeline.from_pretrained("ostris/OpenFLUX.1", torch_dtype=dtype, num_single_layers=0, chunk_size=0).to("cuda")
-pipe.config = FluxTransformer2DModel.from_pretrained("ostris/OpenFLUX.1", subfolder="transformer", num_single_layers=0, chunk_size=0)
+model = FluxTransformer2DModel.from_pretrained("ostris/OpenFLUX.1", torch_dtype=dtype).to("cuda")
+model.num_single_layers="0"
+model.chunk_size="0"
+model.pooled_projections="(_, _, 1)[0]"
+model.pooled_projections_dim="1"
+pipe = FluxWithCFGPipeline.from_pretrained("ostris/OpenFLUX.1", torch_dtype=dtype).to("cuda")
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")
 
 pipe.to("cuda")
@@ -57,9 +61,6 @@ pipe.text_encoder = clip_model.text_model
 pipe.text_encoder.dtype = torch.bfloat16
 torch.cuda.empty_cache()
 
-num_single_layers=0 
-chunk_size=0
-
 MAX_SEED = 2**32-1
 
 class calculateDuration: