Update app2.py

app2.py

@@ -5,8 +5,12 @@ import logging
 import torch
 from PIL import Image
 from os import path
+from torchvision import transforms
+from dataclasses import dataclass
+import math
+from typing import Callable
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
 from diffusers.models.transformers import FluxTransformer2DModel
@@ -14,6 +18,7 @@ import copy
 import random
 import time
 import safetensors.torch
+from tqdm import tqdm
 from safetensors.torch import load_file
 from huggingface_hub import HfFileSystem, ModelCard
 from huggingface_hub import login, hf_hub_download
@@ -34,7 +39,7 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "AlekseyCalvin/SilverAgePoets_FluxS_TestAlpha_Diffusers"
+base_model = "AlekseyCalvin/Artsy_Lite_Flux_v1_by_jurdn_Diffusers"
 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()
@@ -50,12 +55,13 @@ if clipmodel == "norm":
     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)
+t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
 
 pipe.tokenizer = clip_processor.tokenizer
 pipe.text_encoder = clip_model.text_model
 pipe.tokenizer_max_length = maxtokens
 pipe.text_encoder.dtype = torch.bfloat16
-
+pipe.text_encoder_2 = t5.text_model
 
 MAX_SEED = 2**32-1