Update app.py

app.py

@@ -6,6 +6,10 @@ from diffusers import DiffusionPipeline
 import random
 from huggingface_hub import login
 import os
+from diffusers import FluxTransformer2DModel, FluxPipeline
+from transformers import T5EncoderModel, CLIPTextModel
+from optimum.quanto import freeze, qfloat8, quantize
+
 
 
 torch.backends.cudnn.deterministic = True
@@ -17,9 +21,21 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 
 
-# Initialize the base model and specific LoRA
-model_url = "https://huggingface.co/SG161222/RealFlux_1.0b_Dev/resolve/main/1%20-%20Compact%20Version/RealFlux_1.0b_Dev_Compact.safetensors"
-pipe = DiffusionPipeline.from_pretrained("SG161222/RealFlux_1.0b_Dev", torch_dtype=torch.bfloat16)
+bfl_repo = "black-forest-labs/FLUX.1-dev"
+dtype = torch.bfloat16
+
+transformer = FluxTransformer2DModel.from_single_file("https://huggingface.co/SG161222/RealFlux_1.0b_Dev/resolve/main/1%20-%20Compact%20Version/RealFlux_1.0b_Dev_Compact.safetensors", torch_dtype=dtype)
+quantize(transformer, weights=qfloat8)
+freeze(transformer)
+
+text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
+quantize(text_encoder_2, weights=qfloat8)
+freeze(text_encoder_2)
+
+pipe = FluxPipeline.from_pretrained(bfl_repo, transformer=None, text_encoder_2=None, torch_dtype=dtype)
+pipe.transformer = transformer
+pipe.text_encoder_2 = text_encoder_2
+
 
 pipe.to("cuda")