Update app.py

app.py

@@ -20,9 +20,7 @@ from diffusers import AutoencoderKL
 from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
-from safetensors.torch import load_file
-# from torch.profiler import profile, record_function, ProfilerActivity
-# from optimum.quanto import freeze, qfloat8, quantize
+from optimum.quanto import freeze, qfloat8, quantize
 
 
 # ---------------- Encoders ----------------
@@ -67,8 +65,8 @@ device = "cuda"
 t5 = HFEmbedder("google/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
 clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
 ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
-# quantize(t5, weights=qfloat8)
-# freeze(t5)
+quantize(t5, weights=qfloat8)
+freeze(t5)
 
 
 # ---------------- NF4 ----------------
@@ -746,6 +744,7 @@ sd = load_file(hf_hub_download(repo_id="lllyasviel/flux1-dev-bnb-nf4", filename=
 sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items() if "model.diffusion_model" in k}
 model = Flux().to(dtype=torch.bfloat16, device="cuda")
 result = model.load_state_dict(sd)
+model_zero_init = False
 print(result)
 
 # model = Flux().to(dtype=torch.bfloat16, device="cuda")
@@ -764,8 +763,10 @@ def generate_image(
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch_device = torch.device(device)
     
-    global model
-    model = model.to(torch_device)
+    global model, model_zero_init
+    if not model_zero_init:
+        model = model.to(torch_device)
+        model_zero_init = True
 
     if do_img2img and init_image is not None:
         init_image = get_image(init_image)