Update app.py

app.py

@@ -5,16 +5,24 @@ from diffusers import AutoPipelineForText2Image
 from io import BytesIO
 import gradio as gr
 from concurrent.futures import ProcessPoolExecutor
+import multiprocessing
 
 # Load the model once outside of the function
 print("Loading the model...")
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 print("Model loaded successfully.")
 
-def generate_image(prompt, prompt_name):
+def initialize_model():
+    global model
+    if 'model' not in globals():
+        print("Loading the model in worker...")
+        model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
+        print("Model loaded in worker successfully.")
+
+def generate_image(prompt, prompt_name, num_inference_steps=1):
     try:
         print(f"Generating response for {prompt_name} with prompt: {prompt}")
-        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        output = model(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0)
         print(f"Output for {prompt_name}: {output}")
 
         # Check if the model returned images
@@ -47,11 +55,11 @@ async def queue_api_calls(sentence_mapping, character_dict, selected_style):
         prompts.append((paragraph_number, prompt))
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
-    # Set max_workers to the total number of prompts
-    max_workers = len(prompts)
+    # Set max_workers to the minimum of the number of prompts and available CPU cores
+    max_workers = min(len(prompts), multiprocessing.cpu_count())
 
     # Generate images for each prompt in parallel using multiprocessing
-    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+    with ProcessPoolExecutor(max_workers=max_workers, initializer=initialize_model) as executor:
         loop = asyncio.get_running_loop()
         tasks = [loop.run_in_executor(executor, generate_image, prompt, f"Prompt {paragraph_number}") for paragraph_number, prompt in prompts]
         print("Tasks created for image generation.")