Update app.py

app.py

@@ -15,11 +15,13 @@ from pytorch_lightning import seed_everything
 from omegaconf import OmegaConf
 from einops import rearrange, repeat
 from tqdm import tqdm
-from diffusers import FluxPipeline, DiffusionPipeline, EulerAncestralDiscreteScheduler
+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
 from src.utils.camera_util import (
@@ -69,11 +71,24 @@ if cuda_path:
 else:
     print("CUDA installation not found")
 
-# Load Flux pipeline
-flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=huggingface_token)
-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)
+
+print('Loading and fusing lora, please wait...')
+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
+pipe.fuse_lora(lora_scale=1.)
+pipe.unload_lora_weights()
+
+print('Quantizing, please wait...')
+quantize(pipe.transformer, qfloat8)
+freeze(pipe.transformer)
+print('Model quantized!')
+pipe.to('cuda')
+
 
 # Load 3D generation models
 config_path = 'configs/instant-mesh-large.yaml'
@@ -135,18 +150,20 @@ def preprocess(input_image, do_remove_background):
         input_image = resize_foreground(input_image, 0.85)
     return input_image
 
+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
 def generate_mvs(input_image, sample_steps, sample_seed):