Merge pull request #340 from mit-han-lab/dev

feat: support PuLID, Double FBCache and TeaCache; better linter

app/flux.1/fill/assets/description.html

 <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
     <div>
         <h1>
-            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/logo.svg"
+            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
                  alt="logo"
                  style="height: 40px; width: auto; display: block; margin: auto;"/>
             INT4 FLUX.1-fill-dev Demo

app/flux.1/fill/run_gradio.py

@@ -8,25 +8,25 @@ import GPUtil
 import torch
 from diffusers import FluxFillPipeline
 from PIL import Image
+from utils import get_args
+from vars import DEFAULT_GUIDANCE, DEFAULT_INFERENCE_STEP, DEFAULT_STYLE_NAME, EXAMPLES, MAX_SEED, STYLE_NAMES, STYLES
 
 from nunchaku.models.safety_checker import SafetyChecker
 from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
-from utils import get_args
-from vars import DEFAULT_GUIDANCE, DEFAULT_INFERENCE_STEP, DEFAULT_STYLE_NAME, EXAMPLES, MAX_SEED, STYLE_NAMES, STYLES
 
 # import gradio last to avoid conflicts with other imports
-import gradio as gr
+import gradio as gr  # noqa: isort: skip
 
 args = get_args()
 
 if args.precision == "bf16":
-    pipeline = FluxFillPipeline.from_pretrained(f"black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16)
+    pipeline = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16)
     pipeline = pipeline.to("cuda")
     pipeline.precision = "bf16"
 else:
     assert args.precision == "int4"
     pipeline_init_kwargs = {}
-    transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-int4-flux.1-fill-dev")
+    transformer = NunchakuFluxTransformer2dModel.from_pretrained("mit-han-lab/svdq-int4-flux.1-fill-dev")
     pipeline_init_kwargs["transformer"] = transformer
     if args.use_qencoder:
         from nunchaku.models.text_encoders.t5_encoder import NunchakuT5EncoderModel
@@ -35,7 +35,7 @@ else:
         pipeline_init_kwargs["text_encoder_2"] = text_encoder_2
 
     pipeline = FluxFillPipeline.from_pretrained(
-        f"black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16, **pipeline_init_kwargs
+        "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16, **pipeline_init_kwargs
     )
     pipeline = pipeline.to("cuda")
     pipeline.precision = "int4"
@@ -94,7 +94,7 @@ def run(
     return result_image, latency_str
 
 
-with gr.Blocks(css_paths="assets/style.css", title=f"SVDQuant Flux.1-Fill-dev Sketch-to-Image Demo") as demo:
+with gr.Blocks(css_paths="assets/style.css", title="SVDQuant Flux.1-Fill-dev Sketch-to-Image Demo") as demo:
     with open("assets/description.html", "r") as f:
         DESCRIPTION = f.read()
     gpus = GPUtil.getGPUs()
@@ -104,7 +104,7 @@ with gr.Blocks(css_paths="assets/style.css", title=f"SVDQuant Flux.1-Fill-dev Sk
         device_info = f"Running on {gpu.name} with {memory:.0f} GiB memory."
     else:
         device_info = "Running on CPU 🥶 This demo does not work on CPU."
-    notice = f'<strong>Notice:</strong>&nbsp;We will replace unsafe prompts with a default prompt: "A peaceful world."'
+    notice = '<strong>Notice:</strong>&nbsp;We will replace unsafe prompts with a default prompt: "A peaceful world."'
 
     def get_header_str():
 

app/flux.1/redux/assets/description.html

 <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
     <div>
         <h1>
-            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/logo.svg"
+            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
                  alt="logo"
                  style="height: 40px; width: auto; display: block; margin: auto;"/>
             INT4 FLUX.1-redux-dev Demo

app/flux.1/redux/run_gradio.py

@@ -5,16 +5,16 @@ import time
 from datetime import datetime
 
 import GPUtil
+
+# import gradio last to avoid conflicts with other imports
+import gradio as gr
 import torch
 from diffusers import FluxPipeline, FluxPriorReduxPipeline
 from PIL import Image
-
-from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
 from utils import get_args
 from vars import DEFAULT_GUIDANCE, DEFAULT_INFERENCE_STEP, EXAMPLES, MAX_SEED
 
-# import gradio last to avoid conflicts with other imports
-import gradio as gr
+from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
 
 args = get_args()
 
@@ -76,7 +76,7 @@ def run(image, num_inference_steps: int, guidance_scale: float, seed: int) -> tu
     return result_image, latency_str
 
 
-with gr.Blocks(css_paths="assets/style.css", title=f"SVDQuant Flux.1-redux-dev Demo") as demo:
+with gr.Blocks(css_paths="assets/style.css", title="SVDQuant Flux.1-redux-dev Demo") as demo:
     with open("assets/description.html", "r") as f:
         DESCRIPTION = f.read()
     gpus = GPUtil.getGPUs()

app/flux.1/redux/utils.py

@@ -11,9 +11,7 @@ def get_args() -> argparse.Namespace:
         choices=["int4", "bf16"],
         help="Which precisions to use",
     )
-    parser.add_argument(
-        "--count-use", action="store_true", help="Whether to count the number of uses"
-    )
+    parser.add_argument("--count-use", action="store_true", help="Whether to count the number of uses")
     parser.add_argument("--gradio-root-path", type=str, default="")
     args = parser.parse_args()
     return args

app/flux.1/sketch/assets/description.html

 <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
     <div>
         <h1>
-            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/logo.svg"
+            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
                  alt="logo"
                  style="height: 40px; width: auto; display: block; margin: auto;"/>
             INT4 FLUX.1-schnell Sketch-to-Image Demo

app/flux.1/sketch/run.py

 import argparse
 
 import torch
-
 from flux_pix2pix_pipeline import FluxPix2pixTurboPipeline
 
 

app/flux.1/sketch/run_gradio.py

@@ -8,16 +8,16 @@ from datetime import datetime
 import GPUtil
 import numpy as np
 import torch
+from flux_pix2pix_pipeline import FluxPix2pixTurboPipeline
 from PIL import Image
+from utils import get_args
+from vars import DEFAULT_SKETCH_GUIDANCE, DEFAULT_STYLE_NAME, MAX_SEED, STYLE_NAMES, STYLES
 
-from flux_pix2pix_pipeline import FluxPix2pixTurboPipeline
 from nunchaku.models.safety_checker import SafetyChecker
 from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
-from utils import get_args
-from vars import DEFAULT_SKETCH_GUIDANCE, DEFAULT_STYLE_NAME, MAX_SEED, STYLE_NAMES, STYLES
 
 # import gradio last to avoid conflicts with other imports
-import gradio as gr
+import gradio as gr  # noqa: isort: skip
 
 blank_image = Image.new("RGB", (1024, 1024), (255, 255, 255))
 
@@ -109,7 +109,7 @@ def run(image, prompt: str, prompt_template: str, sketch_guidance: float, seed:
     return result_image, latency_str
 
 
-with gr.Blocks(css_paths="assets/style.css", title=f"SVDQuant Sketch-to-Image Demo") as demo:
+with gr.Blocks(css_paths="assets/style.css", title="SVDQuant Sketch-to-Image Demo") as demo:
     with open("assets/description.html", "r") as f:
         DESCRIPTION = f.read()
     gpus = GPUtil.getGPUs()
@@ -119,7 +119,7 @@ with gr.Blocks(css_paths="assets/style.css", title=f"SVDQuant Sketch-to-Image De
         device_info = f"Running on {gpu.name} with {memory:.0f} GiB memory."
     else:
         device_info = "Running on CPU 🥶 This demo does not work on CPU."
-    notice = f'<strong>Notice:</strong>&nbsp;We will replace unsafe prompts with a default prompt: "A peaceful world."'
+    notice = '<strong>Notice:</strong>&nbsp;We will replace unsafe prompts with a default prompt: "A peaceful world."'
 
     def get_header_str():
 

app/flux.1/t2i/assets/description.html

 <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
     <div>
         <h1>
-            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/logo.svg"
+            <img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
                  alt="logo"
                  style="height: 40px; width: auto; display: block; margin: auto;"/>
             FLUX.1-{model} Demo

app/flux.1/t2i/evaluate.py

@@ -2,9 +2,8 @@ import argparse
 import os
 
 import torch
-from tqdm import tqdm
-
 from data import get_dataset
+from tqdm import tqdm
 from utils import get_pipeline, hash_str_to_int
 
 

app/flux.1/t2i/generate.py

@@ -2,7 +2,6 @@ import argparse
 import os
 
 import torch
-
 from utils import get_pipeline
 from vars import PROMPT_TEMPLATES
 

app/flux.1/t2i/latency.py

@@ -4,7 +4,6 @@ import time
 import torch
 from torch import nn
 from tqdm import trange
-
 from utils import get_pipeline
 
 

app/flux.1/t2i/metrics/image_reward.py

 import os
 
-import ImageReward as RM
 import datasets
+import ImageReward as RM
 import torch
 from tqdm import tqdm