Update app.py

app.py

@@ -68,9 +68,17 @@ repo_id = "black-forest-labs/FLUX.1-dev" # Verify if it's 1 or 2 on the official
 hf_token = os.environ.get("HF_TOKEN")
 
 # Load the transformer with the token
+nf4_config = DiffBitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
 dit = Flux2Transformer2DModel.from_pretrained(
     "black-forest-labs/FLUX.2-dev",
     subfolder="transformer",
+    quantization_config=nf4_config,
+    torch_dtype=torch.bfloat16,
     token=os.environ.get("HF_TOKEN"),
 )