fix: add timeout to prevent unexcepted tasl

app.py

@@ -60,6 +60,7 @@ samplers_k_diffusion = [
 # ]
 
 start_time = time.time()
+timeout = 120
 
 scheduler = DDIMScheduler.from_pretrained(
     base_model,
@@ -257,6 +258,8 @@ def inference(
         "sampler_opt": sampler_opt,
         "pww_state": state,
         "pww_attn_weight": g_strength,
+        "start_time": start_time,
+        "timeout": timeout,
     }
 
     if img_input is not None:

modules/model.py

@@ -6,6 +6,7 @@ import re
 from collections import defaultdict
 from typing import List, Optional, Union
 
+import time
 import k_diffusion
 import numpy as np
 import PIL
@@ -446,6 +447,8 @@ class StableDiffusionPipeline(DiffusionPipeline):
         pww_attn_weight=1.0,
         sampler_name="",
         sampler_opt={},
+        start_time=-1,
+        timeout=180,
         scale_ratio=8.0,
     ):
         sampler = self.get_scheduler(sampler_name)
@@ -504,6 +507,9 @@ class StableDiffusionPipeline(DiffusionPipeline):
 
         def model_fn(x, sigma):
 
+            if start_time > 0 and timeout > 0:
+                assert (time.time() - start_time) < timeout, "inference process timed out"
+
             latent_model_input = torch.cat([x] * 2)
             weight_func = lambda w, sigma, qk: w * math.log(1 + sigma) * qk.max()
             encoder_state = {
@@ -617,6 +623,8 @@ class StableDiffusionPipeline(DiffusionPipeline):
         pww_attn_weight=1.0,
         sampler_name="",
         sampler_opt={},
+        start_time=-1,
+        timeout=180,
     ):
         sampler = self.get_scheduler(sampler_name)
         # 1. Check inputs. Raise error if not correct
@@ -667,6 +675,9 @@ class StableDiffusionPipeline(DiffusionPipeline):
 
         def model_fn(x, sigma):
 
+            if start_time > 0 and timeout > 0:
+                assert (time.time() - start_time) < timeout, "inference process timed out"
+
             latent_model_input = torch.cat([x] * 2)
             weight_func = lambda w, sigma, qk: w * math.log(1 + sigma) * qk.max()
             encoder_state = {