improve warmup
Browse files
wfx.py
CHANGED
@@ -98,7 +98,7 @@ class WFX():
|
|
98 |
|
99 |
def warmup(self) -> None:
|
100 |
warmed = 0
|
101 |
-
warmed_total =
|
102 |
|
103 |
warmup_kwargs = dict(
|
104 |
prompt='a photo of a cat',
|
@@ -112,10 +112,17 @@ class WFX():
|
|
112 |
logger.info(f'warming up T2I pipeline...')
|
113 |
for _ in range(warmed_total):
|
114 |
begin = torch.cuda.Event(enable_timing=True)
|
115 |
-
self.T2IPipeline(**warmup_kwargs)
|
116 |
end = torch.cuda.Event(enable_timing=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
117 |
warmed += 1
|
118 |
-
logger.info(f'warmed {warmed}/{warmed_total} - {
|
119 |
|
120 |
|
121 |
|
|
|
98 |
|
99 |
def warmup(self) -> None:
|
100 |
warmed = 0
|
101 |
+
warmed_total = 5
|
102 |
|
103 |
warmup_kwargs = dict(
|
104 |
prompt='a photo of a cat',
|
|
|
112 |
logger.info(f'warming up T2I pipeline...')
|
113 |
for _ in range(warmed_total):
|
114 |
begin = torch.cuda.Event(enable_timing=True)
|
|
|
115 |
end = torch.cuda.Event(enable_timing=True)
|
116 |
+
|
117 |
+
begin.record()
|
118 |
+
self.T2IPipeline(**warmup_kwargs)
|
119 |
+
end.record()
|
120 |
+
|
121 |
+
torch.cuda.synchronize()
|
122 |
+
elapsed_time = begin.elapsed_time(end)
|
123 |
+
|
124 |
warmed += 1
|
125 |
+
logger.info(f'warmed {warmed}/{warmed_total} - {elapsed_time:.2f}ms')
|
126 |
|
127 |
|
128 |
|