Update handler.py
Browse files- handler.py +6 -6
handler.py
CHANGED
|
@@ -8,7 +8,7 @@ from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
|
| 8 |
import time
|
| 9 |
import uuid
|
| 10 |
from huggingface_hub import HfApi
|
| 11 |
-
|
| 12 |
class EndpointHandler:
|
| 13 |
def __init__(self, path=""):
|
| 14 |
self.pipe = FluxPipeline.from_pretrained(
|
|
@@ -22,9 +22,9 @@ class EndpointHandler:
|
|
| 22 |
self.pipe.vae = torch.compile(
|
| 23 |
self.pipe.vae, mode="max-autotune-no-cudagraphs",
|
| 24 |
)
|
| 25 |
-
self.record=0
|
| 26 |
|
| 27 |
def __call__(self, data: Dict[str, Any]) -> str:
|
|
|
|
| 28 |
logger.info(f"Received incoming request with {data=}")
|
| 29 |
|
| 30 |
if "inputs" in data and isinstance(data["inputs"], str):
|
|
@@ -37,7 +37,7 @@ class EndpointHandler:
|
|
| 37 |
" prompt to use for the image generation, and it needs to be a non-empty string."
|
| 38 |
)
|
| 39 |
if prompt=="get_quene":
|
| 40 |
-
return
|
| 41 |
parameters = data.pop("parameters", {})
|
| 42 |
|
| 43 |
num_inference_steps = parameters.get("num_inference_steps", 28)
|
|
@@ -48,7 +48,7 @@ class EndpointHandler:
|
|
| 48 |
# seed generator (seed cannot be provided as is but via a generator)
|
| 49 |
seed = parameters.get("seed", 0)
|
| 50 |
generator = torch.manual_seed(seed)
|
| 51 |
-
|
| 52 |
start_time = time.time()
|
| 53 |
time.sleep(6)
|
| 54 |
# result = self.pipe( # type: ignore
|
|
@@ -62,6 +62,6 @@ class EndpointHandler:
|
|
| 62 |
end_time = time.time()
|
| 63 |
time_taken = end_time - start_time
|
| 64 |
print(f"Time taken: {time_taken:.2f} seconds")
|
| 65 |
-
|
| 66 |
|
| 67 |
-
return
|
|
|
|
| 8 |
import time
|
| 9 |
import uuid
|
| 10 |
from huggingface_hub import HfApi
|
| 11 |
+
record=0
|
| 12 |
class EndpointHandler:
|
| 13 |
def __init__(self, path=""):
|
| 14 |
self.pipe = FluxPipeline.from_pretrained(
|
|
|
|
| 22 |
self.pipe.vae = torch.compile(
|
| 23 |
self.pipe.vae, mode="max-autotune-no-cudagraphs",
|
| 24 |
)
|
|
|
|
| 25 |
|
| 26 |
def __call__(self, data: Dict[str, Any]) -> str:
|
| 27 |
+
global record
|
| 28 |
logger.info(f"Received incoming request with {data=}")
|
| 29 |
|
| 30 |
if "inputs" in data and isinstance(data["inputs"], str):
|
|
|
|
| 37 |
" prompt to use for the image generation, and it needs to be a non-empty string."
|
| 38 |
)
|
| 39 |
if prompt=="get_quene":
|
| 40 |
+
return record
|
| 41 |
parameters = data.pop("parameters", {})
|
| 42 |
|
| 43 |
num_inference_steps = parameters.get("num_inference_steps", 28)
|
|
|
|
| 48 |
# seed generator (seed cannot be provided as is but via a generator)
|
| 49 |
seed = parameters.get("seed", 0)
|
| 50 |
generator = torch.manual_seed(seed)
|
| 51 |
+
record+=1
|
| 52 |
start_time = time.time()
|
| 53 |
time.sleep(6)
|
| 54 |
# result = self.pipe( # type: ignore
|
|
|
|
| 62 |
end_time = time.time()
|
| 63 |
time_taken = end_time - start_time
|
| 64 |
print(f"Time taken: {time_taken:.2f} seconds")
|
| 65 |
+
record-=1
|
| 66 |
|
| 67 |
+
return record
|