Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import torch
 from diffusers import AutoPipelineForText2Image
 from io import BytesIO
 from generate_propmts import generate_prompt
@@ -9,7 +10,7 @@ model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
 async def generate_image(prompt):
     try:
-        # Truncate prompt if necessary
+        # Generate an image based on the prompt
         output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
 
@@ -22,7 +23,6 @@ async def generate_image(prompt):
             return image_bytes
         else:
             raise Exception("No images returned by the model.")
-            
     except Exception as e:
         print(f"Error generating image: {e}")
         return None
@@ -39,14 +39,17 @@ async def process_prompt(sentence_mapping, character_dict, selected_style):
         prompts.append((paragraph_number, prompt))
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
+    # Create tasks for all prompts and run them concurrently
     tasks = [generate_image(prompt) for _, prompt in prompts]
+    results = await asyncio.gather(*tasks)
 
-    for (paragraph_number, _), image_bytes in zip(prompts, await asyncio.gather(*tasks)):
-        if image_bytes:
-            images[paragraph_number] = image_bytes
+    # Map results back to paragraphs
+    for i, (paragraph_number, _) in enumerate(prompts):
+        images[paragraph_number] = results[i]
 
     return images
 
+# Gradio interface with high concurrency limit
 gradio_interface = gr.Interface(
     fn=process_prompt,
     inputs=[
@@ -54,9 +57,9 @@ gradio_interface = gr.Interface(
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
-    outputs="json",
-    concurrency_limit=10  # Allow up to 10 concurrent executions
+    outputs="text",
+    concurrency_limit=20  # Set a high concurrency limit
 )
 
 if __name__ == "__main__":
-    gradio_interface.launch()
+    gradio_interface.launch(share=True)  # Optionally make it public with share=True