quantized

app.py

@@ -19,6 +19,8 @@ from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 import gradio as gr
 import shutil
 import tempfile
+from functools import partial
+from optimum.quanto import quantize, qfloat8, freeze
 from flux_8bit_lora import FluxPipeline
 
 from src.utils.train_util import instantiate_from_config
@@ -72,21 +74,26 @@ else:
 
 device = torch.device('cuda')
 
-# Load Flux pipeline
-flux_pipe = FluxPipeline.from_pretrained(
-    "Freepik/flux.1-lite-8B-alpha",
-    torch_dtype=torch.bfloat16
-)
-flux_pipe.load_lora_weights(hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.safetensors"))
-flux_pipe.fuse_lora(lora_scale=1)
-flux_pipe.to(device="cuda", dtype=torch.bfloat16)
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16, token=huggingface_token)
 
+# Load and fuse LoRA BEFORE quantizing
 print('Loading and fusing lora, please wait...')
-flux_pipe.load_lora_weights(hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.safetensors"))
-# We need this scaling because SimpleTuner fixes the alpha to 16, might be fixed later in diffusers
-# See https://github.com/huggingface/diffusers/issues/9134
-flux_pipe.fuse_lora(lora_scale=1.)
-flux_pipe.unload_lora_weights()
+lora_path = hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.safetensors")
+pipe.load_lora_weights(lora_path)
+pipe.fuse_lora(lora_scale=1.0)
+pipe.unload_lora_weights()
+pipe.transformer.to(device, dtype=torch.bfloat16)
+
+# Now quantize after LoRA is fused
+print('Quantizing, please wait...')
+# Try qint8 if qfloat8 produces invalid values
+quantize(pipe.transformer, qfloat8)  # Consider changing to qint8 if you get invalid values
+freeze(pipe.transformer)
+pipe.transformer.to(device, dtype=torch.bfloat16)
+print('Model quantized!')
+pipe.enable_model_cpu_offload()
+
 
 # Load 3D generation models
 config_path = 'configs/instant-mesh-large.yaml'
@@ -150,16 +157,15 @@ ts_cutoff = 2
 
 @spaces.GPU
 def generate_flux_image(prompt, height, width, steps, scales, seed):
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("Flux inference"):
-        return flux_pipe(
-            prompt=[prompt],
-            generator=torch.Generator().manual_seed(int(seed)),
-            num_inference_steps=int(steps),
-            guidance_scale=float(scales),
-            height=int(height),
-            width=int(width),
-            max_sequence_length=256
-        ).images[0]
+    return pipe(
+        prompt=prompt, 
+        width=int(height),
+        height=int(width),
+        num_inference_steps=int(steps), 
+        generator=torch.Generator().manual_seed(int(seed)),
+        guidance_scale=float(scales),
+        timestep_to_start_cfg=ts_cutoff,
+    ).images[0]
 
 
 @spaces.GPU