Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import asyncio
+import time
 from generate_prompts import generate_prompt
 from diffusers import AutoPipelineForText2Image
 from io import BytesIO
@@ -16,10 +17,12 @@ class ModelActor:
         print("Model loaded successfully.")
 
     def generate_image(self, prompt, prompt_name):
+        start_time = time.time()
+        process_id = os.getpid()
         try:
-            print(f"Generating response for {prompt_name} with prompt: {prompt}")
+            print(f"[{process_id}] Generating response for {prompt_name} with prompt: {prompt}")
             output = self.model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
-            print(f"Output for {prompt_name}: {output}")
+            print(f"[{process_id}] Output for {prompt_name}: {output}")
 
             if isinstance(output.images, list) and len(output.images) > 0:
                 image = output.images[0]
@@ -27,15 +30,17 @@ class ModelActor:
                 try:
                     image.save(buffered, format="JPEG")
                     image_bytes = buffered.getvalue()
-                    print(f"Image bytes length for {prompt_name}: {len(image_bytes)}")
+                    end_time = time.time()
+                    print(f"[{process_id}] Image bytes length for {prompt_name}: {len(image_bytes)}")
+                    print(f"[{process_id}] Time taken for {prompt_name}: {end_time - start_time} seconds")
                     return image_bytes
                 except Exception as e:
-                    print(f"Error saving image for {prompt_name}: {e}")
+                    print(f"[{process_id}] Error saving image for {prompt_name}: {e}")
                     return None
             else:
-                raise Exception(f"No images returned by the model for {prompt_name}.")
+                raise Exception(f"[{process_id}] No images returned by the model for {prompt_name}.")
         except Exception as e:
-            print(f"Error generating image for {prompt_name}: {e}")
+            print(f"[{process_id}] Error generating image for {prompt_name}: {e}")
             return None
 
 model_actor = ModelActor.remote()
@@ -54,8 +59,7 @@ async def queue_api_calls(sentence_mapping, character_dict, selected_style):
     tasks = [model_actor.generate_image.remote(prompt, f"Prompt {paragraph_number}") for paragraph_number, prompt in prompts]
     print("Tasks created for image generation.")
 
-    # Use asyncio to gather the results asynchronously
-    responses = await asyncio.gather(*[ray.get(task) for task in tasks])
+    responses = await asyncio.gather(*[asyncio.to_thread(ray.get, task) for task in tasks])
     print("Responses received from image generation tasks.")
 
     images = {paragraph_number: response for (paragraph_number, _), response in zip(prompts, responses)}