Spaces:
Sleeping
Sleeping
more updates.
Browse files
app.py
CHANGED
@@ -56,8 +56,11 @@ pipeline_mapping = {
|
|
56 |
"stabilityai/stable-diffusion-xl-base-1.0",
|
57 |
"TencentARC/t2i-adapter-canny-sdxl-1.0",
|
58 |
),
|
59 |
-
"Kandinsky 2.2 (T2I)": (
|
60 |
-
|
|
|
|
|
|
|
61 |
}
|
62 |
|
63 |
|
@@ -67,6 +70,7 @@ def load_pipeline(
|
|
67 |
do_torch_compile: bool = False,
|
68 |
):
|
69 |
# Get pipeline details.
|
|
|
70 |
pipeline_details = pipeline_mapping[pipeline_to_benchmark]
|
71 |
pipeline_cls = pipeline_details[0]
|
72 |
pipeline_ckpt = pipeline_details[1]
|
@@ -96,14 +100,15 @@ def load_pipeline(
|
|
96 |
pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, controlnet=controlnet)
|
97 |
elif "Adapters" in pipeline_to_benchmark:
|
98 |
pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, adapter=adapter)
|
99 |
-
|
100 |
pipeline.to(device)
|
101 |
|
102 |
# Optionally set memory layout.
|
103 |
if use_channels_last:
|
|
|
104 |
if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
|
105 |
pipeline.unet.to(memory_format=torch.channels_last)
|
106 |
-
elif pipeline_to_benchmark == "Würstchen (T2I)":
|
107 |
pipeline.prior.to(memory_format=torch.channels_last)
|
108 |
pipeline.decoder.to(memory_format=torch.channels_last)
|
109 |
elif pipeline_to_benchmark == "Kandinsky 2.2 (T2I)":
|
@@ -116,16 +121,23 @@ def load_pipeline(
|
|
116 |
|
117 |
# Optional torch compilation.
|
118 |
if do_torch_compile:
|
|
|
119 |
if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
|
120 |
pipeline.unet = torch.compile(
|
121 |
pipeline.unet, mode="reduce-overhead", fullgraph=True
|
122 |
)
|
123 |
elif pipeline_to_benchmark == "Würstchen (T2I)":
|
124 |
-
pipeline.prior = torch.compile(
|
125 |
-
|
126 |
-
|
127 |
-
pipeline.
|
128 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
129 |
if hasattr(pipeline, "controlnet"):
|
130 |
pipeline.controlnet = torch.compile(
|
131 |
pipeline.controlnet, mode="reduce-overhead", fullgraph=True
|
@@ -135,6 +147,7 @@ def load_pipeline(
|
|
135 |
pipeline.adapter, mode="reduce-overhead", fullgraph=True
|
136 |
)
|
137 |
|
|
|
138 |
return pipeline
|
139 |
|
140 |
|
@@ -144,6 +157,11 @@ def generate(
|
|
144 |
use_channels_last: bool = False,
|
145 |
do_torch_compile: bool = False,
|
146 |
):
|
|
|
|
|
|
|
|
|
|
|
147 |
print("Start...")
|
148 |
print("Torch version", torch.__version__)
|
149 |
print("Torch CUDA version", torch.version.cuda)
|
@@ -189,7 +207,7 @@ with gr.Blocks() as demo:
|
|
189 |
use_channels_last = gr.Checkbox(label="Use `channels_last` memory layout?")
|
190 |
pipeline_to_benchmark = gr.Dropdown(
|
191 |
list(pipeline_mapping.keys()),
|
192 |
-
value=
|
193 |
multiselect=False,
|
194 |
label="Pipeline to benchmark",
|
195 |
)
|
|
|
56 |
"stabilityai/stable-diffusion-xl-base-1.0",
|
57 |
"TencentARC/t2i-adapter-canny-sdxl-1.0",
|
58 |
),
|
59 |
+
"Kandinsky 2.2 (T2I)": (
|
60 |
+
AutoPipelineForText2Image,
|
61 |
+
"kandinsky-community/kandinsky-2-2-decoder",
|
62 |
+
),
|
63 |
+
"Würstchen (T2I)": (AutoPipelineForText2Image, "warp-ai/wuerstchen"),
|
64 |
}
|
65 |
|
66 |
|
|
|
70 |
do_torch_compile: bool = False,
|
71 |
):
|
72 |
# Get pipeline details.
|
73 |
+
print(f"Loading pipeline: {pipeline_to_benchmark}")
|
74 |
pipeline_details = pipeline_mapping[pipeline_to_benchmark]
|
75 |
pipeline_cls = pipeline_details[0]
|
76 |
pipeline_ckpt = pipeline_details[1]
|
|
|
100 |
pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, controlnet=controlnet)
|
101 |
elif "Adapters" in pipeline_to_benchmark:
|
102 |
pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, adapter=adapter)
|
103 |
+
|
104 |
pipeline.to(device)
|
105 |
|
106 |
# Optionally set memory layout.
|
107 |
if use_channels_last:
|
108 |
+
print("Setting memory layout.")
|
109 |
if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
|
110 |
pipeline.unet.to(memory_format=torch.channels_last)
|
111 |
+
elif pipeline_to_benchmark == "Würstchen (T2I)":
|
112 |
pipeline.prior.to(memory_format=torch.channels_last)
|
113 |
pipeline.decoder.to(memory_format=torch.channels_last)
|
114 |
elif pipeline_to_benchmark == "Kandinsky 2.2 (T2I)":
|
|
|
121 |
|
122 |
# Optional torch compilation.
|
123 |
if do_torch_compile:
|
124 |
+
print("Compiling pipeline.")
|
125 |
if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
|
126 |
pipeline.unet = torch.compile(
|
127 |
pipeline.unet, mode="reduce-overhead", fullgraph=True
|
128 |
)
|
129 |
elif pipeline_to_benchmark == "Würstchen (T2I)":
|
130 |
+
pipeline.prior = torch.compile(
|
131 |
+
pipeline.prior, mode="reduce-overhead", fullgraph=True
|
132 |
+
)
|
133 |
+
pipeline.decoder = torch.compile(
|
134 |
+
pipeline.decoder, mode="reduce-overhead", fullgraph=True
|
135 |
+
)
|
136 |
+
elif pipeline_to_benchmark == "Kandinsky 2.2 (T2I)":
|
137 |
+
pipeline.unet = torch.compile(
|
138 |
+
pipeline.unet, mode="reduce-overhead", fullgraph=True
|
139 |
+
)
|
140 |
+
|
141 |
if hasattr(pipeline, "controlnet"):
|
142 |
pipeline.controlnet = torch.compile(
|
143 |
pipeline.controlnet, mode="reduce-overhead", fullgraph=True
|
|
|
147 |
pipeline.adapter, mode="reduce-overhead", fullgraph=True
|
148 |
)
|
149 |
|
150 |
+
print("Pipeline loaded.")
|
151 |
return pipeline
|
152 |
|
153 |
|
|
|
157 |
use_channels_last: bool = False,
|
158 |
do_torch_compile: bool = False,
|
159 |
):
|
160 |
+
if isinstance(pipeline_to_benchmark, list):
|
161 |
+
# It can only happen when we don't select a pipeline to benchmark.
|
162 |
+
raise ValueError(
|
163 |
+
"pipeline_to_benchmark cannot be None. Please select a pipeline to benchmark."
|
164 |
+
)
|
165 |
print("Start...")
|
166 |
print("Torch version", torch.__version__)
|
167 |
print("Torch CUDA version", torch.version.cuda)
|
|
|
207 |
use_channels_last = gr.Checkbox(label="Use `channels_last` memory layout?")
|
208 |
pipeline_to_benchmark = gr.Dropdown(
|
209 |
list(pipeline_mapping.keys()),
|
210 |
+
value=None,
|
211 |
multiselect=False,
|
212 |
label="Pipeline to benchmark",
|
213 |
)
|