Skip to content

Commit e419dc2

Browse files
lmxyylmxyy
authored
fix: update the kontext demos (#515)
* fix the use count * update * change the example step to 20 * allow embedding website * support fp4 on blackwell * fp4 kontext runnable * add more examples * add kontext examples
1 parent 121ee75 commit e419dc2

11 files changed

Lines changed: 73 additions & 57 deletions

File tree

‎app/flux.1/depth_canny/assets/description.html‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,13 @@
22
<div>
33
<!-- Logo Row -->
44
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-bottom: 10px;">
5-
<a href="https://github.com/mit-han-lab/nunchaku">
5+
<a href="https://github.com/mit-han-lab/nunchaku" target="_blank">
66
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/nunchaku.svg"
7-
alt="nunchaku logo"
8-
style="height: 150px; width: auto;"/>
7+
alt="nunchaku logo" style="height: 150px; width: auto;" />
98
</a>
10-
<a href="https://hanlab.mit.edu/projects/svdquant">
9+
<a href="https://hanlab.mit.edu/projects/svdquant" target="_blank">
1110
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
12-
alt="svdquant logo"
13-
style="height: 40px; width: auto;"/>
11+
alt="svdquant logo" style="height: 40px; width: auto;" />
1412
</a>
1513
</div>
1614
<h1 style="margin-top: 0;">INT4 FLUX.1-{model_name}-dev Demo</h1>

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

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,18 @@
22
<div>
33
<!-- Logo Row -->
44
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-bottom: 10px;">
5-
<a href="https://github.com/mit-han-lab/nunchaku">
5+
<a href="https://github.com/mit-han-lab/nunchaku" target="_blank">
66
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/nunchaku.svg"
7-
alt="nunchaku logo"
8-
style="height: 150px; width: auto;"/>
7+
alt="nunchaku logo" style="height: 150px; width: auto;" />
98
</a>
10-
<a href="https://hanlab.mit.edu/projects/svdquant">
9+
<a href="https://hanlab.mit.edu/projects/svdquant" target="_blank">
1110
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
12-
alt="svdquant logo"
13-
style="height: 40px; width: auto;"/>
11+
alt="svdquant logo" style="height: 40px; width: auto;" />
1412
</a>
1513
</div>
1614

1715
<!-- Title -->
1816
<h1 style="margin-top: 0;">INT4 FLUX.1-fill-dev Demo</h1>
19-
<h4>Quantization Library:
20-
<a href='https://github.com/mit-han-lab/deepcompressor' target="_blank">DeepCompressor</a>&nbsp;
21-
Inference Engine: <a href='https://github.com/mit-han-lab/nunchaku' target="_blank">Nunchaku</a>&nbsp;
22-
</h4>
2317
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
2418
{device_info}
2519
</div>

‎app/flux.1/kontext/assets/description.html‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,23 +2,20 @@
22
<div>
33
<!-- Logo Row -->
44
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-bottom: 10px;">
5-
<a href="https://github.com/mit-han-lab/nunchaku">
5+
<a href="https://github.com/mit-han-lab/nunchaku" target="_blank">
66
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/nunchaku.svg"
77
alt="nunchaku logo" style="height: 150px; width: auto;" />
88
</a>
9-
<a href="https://hanlab.mit.edu/projects/svdquant">
9+
<a href="https://hanlab.mit.edu/projects/svdquant" target="_blank">
1010
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
1111
alt="svdquant logo" style="height: 40px; width: auto;" />
1212
</a>
1313
</div>
14-
<h1 style="margin-top: 0;">INT4 FLUX.1-Kontext-dev Demo</h1>
14+
<h1 style="margin-top: 0;">{precision} FLUX.1-Kontext-dev Demo</h1>
1515

1616
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
1717
{device_info}
1818
</div>
19-
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
20-
{notice}
21-
</div>
2219
{count_info}
2320
</div>
2421
</div>

‎app/flux.1/kontext/run_gradio.py‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,10 @@
2222
pipeline = pipeline.to("cuda")
2323
pipeline.precision = "bf16"
2424
else:
25-
assert args.precision == "int4"
25+
assert args.precision in ["int4", "fp4"]
2626
pipeline_init_kwargs = {}
2727
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
28-
"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-int4_r32-flux.1-kontext-dev.safetensors"
28+
f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{args.precision}_r32-flux.1-kontext-dev.safetensors"
2929
)
3030
pipeline_init_kwargs["transformer"] = transformer
3131
if args.use_qencoder:
@@ -40,7 +40,7 @@
4040
"black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16, **pipeline_init_kwargs
4141
)
4242
pipeline = pipeline.to("cuda")
43-
pipeline.precision = "int4"
43+
pipeline.precision = args.precision
4444

4545

4646
def run(image, prompt: str, num_inference_steps: int, guidance_scale: float, seed: int) -> tuple[Image, str]:
@@ -65,17 +65,17 @@ def run(image, prompt: str, num_inference_steps: int, guidance_scale: float, see
6565
latency_str = f"{latency:.2f}s"
6666
torch.cuda.empty_cache()
6767
if args.count_use:
68-
if os.path.exists(f"{args.model}-use_count.txt"):
69-
with open(f"{args.model}-use_count.txt", "r") as f:
68+
if os.path.exists("use_count.txt"):
69+
with open("use_count.txt", "r") as f:
7070
count = int(f.read())
7171
else:
7272
count = 0
7373
count += 1
7474
current_time = datetime.now()
7575
print(f"{current_time}: {count}")
76-
with open(f"{args.model}-use_count.txt", "w") as f:
76+
with open("use_count.txt", "w") as f:
7777
f.write(str(count))
78-
with open(f"{args.model}-use_record.txt", "a") as f:
78+
with open("use_record.txt", "a") as f:
7979
f.write(f"{current_time}: {count}\n")
8080
return result_image, latency_str
8181

@@ -91,7 +91,6 @@ def run(image, prompt: str, num_inference_steps: int, guidance_scale: float, see
9191
device_info = f"Running on {gpu_name} with {gpu_memory:.0f} GiB memory."
9292
else:
9393
device_info = "Running on CPU 🥶 This demo does not work on CPU."
94-
notice = '<strong>Notice:</strong>&nbsp;We will replace unsafe prompts with a default prompt: "A peaceful world."'
9594

9695
def get_header_str():
9796

@@ -108,7 +107,9 @@ def get_header_str():
108107
)
109108
else:
110109
count_info = ""
111-
header_str = DESCRIPTION.format(device_info=device_info, notice=notice, count_info=count_info)
110+
header_str = DESCRIPTION.format(
111+
precision=args.precision.upper(), device_info=device_info, count_info=count_info
112+
)
112113
return header_str
113114

114115
header = gr.HTML(get_header_str())

‎app/flux.1/kontext/utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
def get_args() -> argparse.Namespace:
55
parser = argparse.ArgumentParser()
66
parser.add_argument(
7-
"-p", "--precision", type=str, default="int4", choices=["int4", "bf16"], help="Which precisions to use"
7+
"-p", "--precision", type=str, default="int4", choices=["int4", "fp4", "bf16"], help="Which precisions to use"
88
)
99
parser.add_argument("--use-qencoder", action="store_true", help="Whether to use 4-bit text encoder")
1010
parser.add_argument("--no-safety-checker", action="store_true", help="Disable safety checker")

‎app/flux.1/kontext/vars.py‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,43 @@
11
MAX_SEED = 1000000000
22

33
EXAMPLES = [
4+
[
5+
"https://images.pexels.com/photos/15460314/pexels-photo-15460314.jpeg",
6+
"Change the color of the woman's dress to red. The background is a beach. "
7+
"The woman is holding a sign that says 'Nunchaku is awesome'",
8+
20,
9+
2.5,
10+
23,
11+
],
12+
[
13+
"https://huggingface.co/mit-han-lab/nunchaku-artifacts/resolve/main/ComfyUI-nunchaku/test_data/logo.png",
14+
"Change the logo of 'MIT HAN Lab' to 'MIT Nunchaku' in the same style.",
15+
20,
16+
2.5,
17+
233,
18+
],
19+
[
20+
"https://huggingface.co/mit-han-lab/nunchaku-artifacts/resolve/main/ComfyUI-nunchaku/test_data/monalisa.jpg",
21+
"Convert the image to ghibli style",
22+
20,
23+
2.5,
24+
2333,
25+
],
26+
[
27+
"https://huggingface.co/mit-han-lab/nunchaku-artifacts/resolve/main/ComfyUI-nunchaku/test_data/mushroom_depth.webp",
28+
"Transform the depth map into an ethereal fantasy image. A glowing mushroom forms the roof of an ancient treehouse, "
29+
"its golden light casting warmth over moss, tiny flowers, and a stone path. Smaller glowing mushrooms create a "
30+
"multi-level home among the tree’s twisted branches. In the background, a misty forest with waterfalls and a "
31+
"starry night sky adds to the magical atmosphere.",
32+
20,
33+
2.5,
34+
23333,
35+
],
436
[
537
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png",
638
"Make Pikachu hold a sign that says 'Nunchaku is awesome', yarn art style, detailed, vibrant colors",
7-
28,
39+
20,
840
2.5,
9-
3,
41+
233333,
1042
],
1143
]

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,18 @@
33

44
<!-- Logo Row -->
55
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-bottom: 10px;">
6-
<a href="https://github.com/mit-han-lab/nunchaku">
6+
<a href="https://github.com/mit-han-lab/nunchaku" target="_blank">
77
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/nunchaku.svg"
8-
alt="nunchaku logo"
9-
style="height: 150px; width: auto;"/>
8+
alt="nunchaku logo" style="height: 150px; width: auto;" />
109
</a>
11-
<a href="https://hanlab.mit.edu/projects/svdquant">
10+
<a href="https://hanlab.mit.edu/projects/svdquant" target="_blank">
1211
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
13-
alt="svdquant logo"
14-
style="height: 40px; width: auto;"/>
12+
alt="svdquant logo" style="height: 40px; width: auto;" />
1513
</a>
1614
</div>
1715

1816
<!-- Title -->
19-
<h1 style="margin-top: 0;">INT4 FLUX.1-redux-dev Demo</h1>
17+
<h1 style="margin-top: 0;">INT4 FLUX.1-Redux-dev Demo</h1>
2018

2119
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
2220
{device_info}

‎app/flux.1/redux/run_gradio.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,10 @@
2626
pipeline = pipeline.to("cuda")
2727
pipeline.precision = "bf16"
2828
else:
29-
assert args.precision == "int4"
29+
assert args.precision in ["int4", "fp4"]
3030
pipeline_init_kwargs = {}
3131
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
32-
"mit-han-lab/nunchaku-flux.1-dev/svdq-int4_r32-flux.1-dev.safetensors"
32+
f"mit-han-lab/nunchaku-flux.1-dev/svdq-{args.precision}_r32-flux.1-dev.safetensors"
3333
)
3434
pipeline = FluxPipeline.from_pretrained(
3535
"black-forest-labs/FLUX.1-dev",
@@ -39,7 +39,7 @@
3939
torch_dtype=torch.bfloat16,
4040
)
4141
pipeline = pipeline.to("cuda")
42-
pipeline.precision = "int4"
42+
pipeline.precision = args.precision
4343

4444

4545
def run(image, num_inference_steps: int, guidance_scale: float, seed: int) -> tuple[Image, str]:

‎app/flux.1/redux/utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ def get_args() -> argparse.Namespace:
88
"--precision",
99
type=str,
1010
default="int4",
11-
choices=["int4", "bf16"],
11+
choices=["int4", "fp4", "bf16"],
1212
help="Which precisions to use",
1313
)
1414
parser.add_argument("--count-use", action="store_true", help="Whether to count the number of uses")

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,13 @@
22
<div>
33
<!-- Logo Row -->
44
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-bottom: 10px;">
5-
<a href="https://github.com/mit-han-lab/nunchaku">
5+
<a href="https://github.com/mit-han-lab/nunchaku" target="_blank">
66
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/nunchaku.svg"
7-
alt="nunchaku logo"
8-
style="height: 150px; width: auto;"/>
7+
alt="nunchaku logo" style="height: 150px; width: auto;" />
98
</a>
10-
<a href="https://hanlab.mit.edu/projects/svdquant">
9+
<a href="https://hanlab.mit.edu/projects/svdquant" target="_blank">
1110
<img src="https://github.com/mit-han-lab/nunchaku/raw/refs/heads/main/assets/svdquant.svg"
12-
alt="svdquant logo"
13-
style="height: 40px; width: auto;"/>
11+
alt="svdquant logo" style="height: 40px; width: auto;" />
1412
</a>
1513
</div>
1614
<h1 style="margin-top: 0;">INT4 FLUX.1-schnell Sketch-to-Image Demo</h1>

0 commit comments

Comments
 (0)