Update app.py

app.py

@@ -1,56 +1,45 @@
 import os
 import asyncio
-import concurrent.futures
+from concurrent.futures import ProcessPoolExecutor
 from io import BytesIO
 from diffusers import StableDiffusionPipeline
 import gradio as gr
 from generate_prompts import generate_prompt
 
-# Initialize model globally
+# Load the model once at the start
+print("Loading the Stable Diffusion model...")
 model = StableDiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo")
+print("Model loaded successfully.")
 
 def generate_image(prompt, prompt_name):
-    """
-    Generates an image based on the provided prompt.
-    Parameters:
-        - prompt (str): The input text for image generation.
-        - prompt_name (str): A name for the prompt, used for logging.
-    Returns:
-        bytes: The generated image data in bytes format, or None if generation fails.
-    """
     try:
         print(f"Generating image for {prompt_name}")
-        output = model(prompt=prompt, num_inference_steps=50, guidance_scale=7.5)
-        if isinstance(output.images, list) and len(output.images) > 0:
+        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        if output and hasattr(output, 'images') and len(output.images) > 0:
+            print(f"Image generated for {prompt_name}")
             image = output.images[0]
             buffered = BytesIO()
             image.save(buffered, format="JPEG")
             image_bytes = buffered.getvalue()
             return image_bytes
         else:
+            print(f"No images found for {prompt_name}")
             return None
     except Exception as e:
         print(f"An error occurred while generating image for {prompt_name}: {e}")
         return None
 
 async def queue_api_calls(sentence_mapping, character_dict, selected_style):
-    """
-    Generates images for all provided prompts in parallel using ProcessPoolExecutor.
-    Parameters:
-        - sentence_mapping (dict): Mapping between paragraph numbers and sentences.
-        - character_dict (dict): Dictionary mapping characters to their descriptions.
-        - selected_style (str): Selected illustration style.
-    Returns:
-        dict: A dictionary where keys are paragraph numbers and values are image data in bytes format.
-    """
+    print("Starting to queue API calls...")
     prompts = []
     for paragraph_number, sentences in sentence_mapping.items():
         combined_sentence = " ".join(sentences)
         prompt = generate_prompt(combined_sentence, sentence_mapping, character_dict, selected_style)
         prompts.append((paragraph_number, prompt))
+        print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
     loop = asyncio.get_running_loop()
-    with concurrent.futures.ProcessPoolExecutor() as pool:
+    with ProcessPoolExecutor() as pool:
         tasks = [
             loop.run_in_executor(pool, generate_image, prompt, f"Prompt {paragraph_number}")
             for paragraph_number, prompt in prompts
@@ -58,32 +47,36 @@ async def queue_api_calls(sentence_mapping, character_dict, selected_style):
         responses = await asyncio.gather(*tasks)
     
     images = {paragraph_number: response for (paragraph_number, _), response in zip(prompts, responses)}
+    print("Finished queuing API calls.")
     return images
 
 def process_prompt(sentence_mapping, character_dict, selected_style):
-    """
-    Processes the provided prompts and generates images.
-    Parameters:
-        - sentence_mapping (dict): Mapping between paragraph numbers and sentences.
-        - character_dict (dict): Dictionary mapping characters to their descriptions.
-        - selected_style (str): Selected illustration style.
-    Returns:
-        dict: A dictionary where keys are paragraph numbers and values are image data in bytes format.
-    """
+    print("Processing prompt...")
+    print(f"Sentence Mapping: {sentence_mapping}")
+    print(f"Character Dict: {character_dict}")
+    print(f"Selected Style: {selected_style}")
     try:
         loop = asyncio.get_running_loop()
+        print("Using existing event loop.")
     except RuntimeError:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
+        print("Created new event loop.")
 
     cmpt_return = loop.run_until_complete(queue_api_calls(sentence_mapping, character_dict, selected_style))
+    print("Prompt processing complete.")
     return cmpt_return
 
 gradio_interface = gr.Interface(
     fn=process_prompt,
-    inputs=[gr.JSON(label="Sentence Mapping"), gr.JSON(label="Character Dict"), gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")],
+    inputs=[
+        gr.JSON(label="Sentence Mapping"),
+        gr.JSON(label="Character Dict"),
+        gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
+    ],
     outputs="json"
 ).queue(default_concurrency_limit=20)  # Set concurrency limit if needed
 
 if __name__ == "__main__":
+    print("Launching Gradio interface...")
     gradio_interface.launch()