Merge branch 'main' into Fix/SubmitInterface

mlcommons · Sep 5, 2024 · e031461 · e031461
2 parents ad47c38 + bf2cfdf
commit e031461
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/backend/app/domain/services/utils/image_generators.py b/backend/app/domain/services/utils/image_generators.py
@@ -205,6 +205,7 @@ def __init__(self, model_instance: str):
         self.api_key = os.getenv("HF")
         self.session = RequestSession()
         self.model_instance = model_instance
+        self.timeout = 62
 
     def generate_images(
         self, prompt: str, num_images: int, model, endpoint, user_id

diff --git a/backend/worker/tasks.py b/backend/worker/tasks.py
@@ -65,6 +65,21 @@ def generate_images(
         res = images.apply_async()
         print(res)
         all_responses = res.get(disable_sync_subtasks=False)
+        successes = len(
+            [response for response in all_responses if response["message"] == "Success"]
+        )
+        if (successes + num_of_current_images) < 5:
+            more_images = celery.group(
+                *[
+                    generate_nibbler_images_celery.s(
+                        prompt, num_images, models, endpoint, user_id
+                    )
+                    for _ in range(5 - successes + num_of_current_images)
+                ]
+            )
+            res = more_images.apply_async()
+            additional_responses = res.get(disable_sync_subtasks=False)
+            all_responses.extend(additional_responses)
         if (len(all_responses) + num_of_current_images) >= 5:
             job_service.remove_registry({"prompt": prompt, "user_id": user_id})
 
@@ -85,5 +100,6 @@ def generate_images(
 
             logger.critical(json.dumps(info_to_log))
     except Exception as e:
-        print(e)
+        error = e
+        print("Error: ", error)
         job_service.remove_registry({"prompt": prompt, "user_id": user_id})