Update app.py

app.py

@@ -24,7 +24,8 @@ import random
 import string
 from diffusers import FluxPipeline
 from huggingface_hub import hf_hub_download
-
+from diffusers.quantizers import PipelineQuantizationConfig
+from diffusers import (FluxPriorReduxPipeline, FluxInpaintPipeline, FluxFillPipeline, FluxKontextPipeline, FluxPipeline)
 
 # Login Hugging Face Hub
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -35,15 +36,16 @@ import diffusers
 dtype = torch.bfloat16
 device = "cuda:0"
 
-print(device)
-#base_model = "black-forest-labs/FLUX.1-dev"
 base_model = "black-forest-labs/FLUX.1-Krea-dev"
-# load pipe
 
-txt2img_pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=dtype)
+pipeline_quant_config = PipelineQuantizationConfig(
+    quant_backend="bitsandbytes_4bit",
+    quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
+    components_to_quantize=["transformer", "text_encoder_2"],
+)
 
+txt2img_pipe = FluxKontextPipeline.from_pretrained(base_model, quantization_config=pipeline_quant_config, torch_dtype=dtype)
 txt2img_pipe = txt2img_pipe.to(device)
-#txt2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
 
 MAX_SEED = 2**32 - 1
 
@@ -157,14 +159,15 @@ def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, s
                     adapter_weights.append(adapter_weight)
                     if lora_repo and weights and adapter_name:
                         try:
-                            #txt2img_pipe.to(device)
                             txt2img_pipe.load_lora_weights(hf_hub_download(lora_repo, weights), adapter_name=lora_name)
                         except:
                             print("load lora error")
+                
                 # set lora weights
-                #if len(lora_names) > 0:
-                    #txt2img_pipe.to(device)
-                    #txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
+                if len(lora_names) > 0:
+                    txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
+                    txt2img_pipe.fuse_lora(adapter_names=lora_names)
+                    txt2img_pipe.enable_vae_slicing()
     
     # Generate image
     error_message = ""