Update app.py

app.py

@@ -13,13 +13,15 @@ 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
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoPipelineForImage2Image, ConfigMixin
 import safetensors.torch
 from safetensors.torch import load_file
 from pipeline import FluxWithCFGPipeline
 from transformers import CLIPModel, CLIPProcessor, CLIPConfig
 import gc
 import warnings
+import safetensors.torch
+
 
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
@@ -35,10 +37,12 @@ with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 dtype = torch.bfloat16
-pipe = FluxWithCFGPipeline.from_pretrained("ostris/OpenFLUX.1", torch_dtype=dtype, text_encoder_3=None, tokenizer_3=None
-).to("cuda")
+
+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")
 clipmodel = 'norm'
 if clipmodel == "long":
@@ -48,11 +52,10 @@ if clipmodel == "long":
 if clipmodel == "norm":
     model_id = "zer0int/CLIP-GmP-ViT-L-14"
     config = CLIPConfig.from_pretrained(model_id)
-    maxtokens = 77
-clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=True).to("cuda")
-clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=True, return_tensors="pt", truncation=True)
-config.text_config.max_position_embeddings = 77
-
+    maxtokens = 512
+clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=False).to("cuda")
+clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=False, return_tensors="pt", truncation=True)
+pipe.transformer=(FluxTransformer2DModel.from_pretrained("ostris/OpenFLUX.1", num_single_layers=0, chunk_size=0)
 pipe.tokenizer = clip_processor.tokenizer
 pipe.text_encoder = clip_model.text_model
 pipe.tokenizer_max_length = maxtokens