Update app.py

app.py

@@ -6,7 +6,10 @@ import torch
 from PIL import Image
 from os import path
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, FlowMatchEulerDiscreteScheduler,
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
 import copy
 import random
 import time
@@ -44,6 +47,23 @@ torch_dtype=torch.bfloat16,
 trust_remote_code=True,
 ).to(device)
 
+clipmodel = 'norm'
+if clipmodel == "long":
+    model_id = "zer0int/LongCLIP-GmP-ViT-L-14"
+    config = CLIPConfig.from_pretrained(model_id)
+    maxtokens = 248
+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)
+
+pipe.tokenizer = clip_processor.tokenizer
+pipe.text_encoder = clip_model.text_model
+pipe.tokenizer_max_length = maxtokens
+pipe.text_encoder.dtype = torch.bfloat16
+
 MAX_SEED = 2**32-1
 
 class calculateDuration: