Update app.py

app.py

@@ -11,16 +11,30 @@ client = Client("https://reverent-atlei2.hf.space/",hf_token=hf_token_id)
 def warmer():
     while True:
         try:
+            start_time = time.time()
             text_output,image_output = client.predict(
 				"atleiwarm0001",	# str  in 'message' Textbox component
 				"",	# str (filepath or URL to image) in 'Upload any Image' Image component
 				api_name="/predict"
             )
+            end_time = time.time()
+            duration = start_time - end_time
+            print("Time elapsed for Pygmalion model : ", duration)
+            if duration > 120: # Check if duration to generate tokens exceeds 2 minutes
+                restart()
+                
+            start_time = time.time()
             text_output_2,image_output_2 = client.predict(
 				"atleiwarm0002",	# str  in 'message' Textbox component
 				"",	# str (filepath or URL to image) in 'Upload any Image' Image component
 				api_name="/predict"
             )
+            end_time = time.time()
+            duration = start_time - end_time
+            print("Time elapsed for WizardLM model : ", duration)
+            if duration > 300: # Check if duration to generate tokens exceeds 5 minutes
+                restart()
+                
             if (text_output != None) and (text_output_2 != None):
                     print("ALIVE")
                     print("pygmalion response : ", text_output)
@@ -33,19 +47,17 @@ def warmer():
             print(str(e))
 
 def restart():
-    while True:
-        time.sleep(10800) # Restart ATLEI2 Space every 3 hours
-        try:
-            text_output,image_output = client.predict(
-				"atleiwarm0003",	# str  in 'message' Textbox component
-				"",	# str (filepath or URL to image) in 'Upload any Image' Image component
-				api_name="/predict"
-            )
-            print("RESTARTED")
+    try:
+        text_output,image_output = client.predict(
+			"atleiwarm0003",	# str  in 'message' Textbox component
+			"",	# str (filepath or URL to image) in 'Upload any Image' Image component
+			api_name="/predict"
+        )
+        print("RESTARTED")
             
-        except Exception as e:
-            print("NO INTERNET")
-            print(str(e))
+    except Exception as e:
+        print("NO INTERNET")
+        print(str(e))
             
 def placeholder(dummy):
     return "lol"