Update app.py

app.py

@@ -3,18 +3,11 @@ import asyncio
 from generate_prompts import generate_prompt
 from diffusers import AutoPipelineForText2Image
 from io import BytesIO
-import json
 import gradio as gr
 
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-
-prompt1 = "write a 5 paragraph explanation of how to use python async and await. Return a JSON structure as follows {'prompt_name': 'prompt1','response': '[response]'}"
-prompt2 = "write a 5 paragraph explanation of limitations for using asyncio.run(). Return a JSON structure as follows {'prompt_name': 'prompt2','response': '[response}'}"
-prompt3 = "write a 5 paragraph explanation of how to use asyncio.get_running_loop(). Return a JSON structure as follows {'prompt_name': 'prompt3','response': '[response]'}"
-prompt4 = "write a 5 paragraph explanation of how to use asyncio.gather(). Return a JSON structure as follows {'prompt_name': 'prompt4','response': '[response]'}"
-
 async def generate_image(prompt, prompt_name):
     try:
         print(f"Generating response for {prompt_name}")
@@ -41,52 +34,44 @@ async def generate_image(prompt, prompt_name):
 async def queue_api_calls(sentence_mapping, character_dict, selected_style):
     print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
     prompts = []
-    
+
     # Generate prompts for each paragraph
     for paragraph_number, sentences in sentence_mapping.items():
         combined_sentence = " ".join(sentences)
         prompt = generate_prompt(combined_sentence, character_dict, selected_style)
         prompts.append((paragraph_number, prompt))
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
-        
+
+    # Generate images for each prompt in parallel
     tasks = [generate_image(prompt, f"Prompt {paragraph_number}") for paragraph_number, prompt in prompts]
-    responses = await asyncio.gather(generate_image(*task))
+    responses = await asyncio.gather(*tasks)
     
-    #Note: Although the API calls get processed in async order, asyncio.gather and returns them in the request order
-	images = {}
-    
-    # Iterate through each response
-    # Map results back to paragraphs
-    for i, (paragraph_number, _) in enumerate(prompts):
-        if i < len(results):
-            images[paragraph_number] = results[i]
-        else:
-            print(f"Error: No result for paragraph {paragraph_number}")
-
+    images = {paragraph_number: response for (paragraph_number, _), response in zip(prompts, responses)}
     return images
 
 def process_prompt(sentence_mapping, character_dict, selected_style):
     try:
-        #see if there is a loop already running. If there is, reuse it.
+        # See if there is a loop already running. If there is, reuse it.
         loop = asyncio.get_running_loop()
     except RuntimeError:
         # Create new event loop if one is not running
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-    #this sends the prompts to function that sets up the async calls. Once all the calls to the API complete, it returns a list of the gr.Textbox with value= set.
+    # This sends the prompts to function that sets up the async calls. Once all the calls to the API complete, it returns a list of the gr.Textbox with value= set.
     cmpt_return = loop.run_until_complete(queue_api_calls(sentence_mapping, character_dict, selected_style))
     return cmpt_return
 
 # Gradio interface with high concurrency limit
 gradio_interface = gr.Interface(
-	fn=process_prompt,
+    fn=process_prompt,
     inputs=[
         gr.JSON(label="Sentence Mapping"),
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
-    outputs="json")
-    
+    outputs="json"
+)
+
 if __name__ == "__main__":
-    demo.launch()
+    gradio_interface.launch()