Update app.py

app.py

@@ -1,9 +1,8 @@
 import os
-import asyncio
-from concurrent.futures import ProcessPoolExecutor
 from io import BytesIO
 from diffusers import AutoPipelineForText2Image
 import gradio as gr
+import base64
 from generate_prompts import generate_prompt
 
 # Load the model once at the start
@@ -17,52 +16,37 @@ def truncate_prompt(prompt, max_length=77):
         prompt = " ".join(tokens[:max_length])
     return prompt
 
-def generate_image(prompt, prompt_name):
+def generate_image(prompt):
     try:
         truncated_prompt = truncate_prompt(prompt)
-        print(f"Generating image for {prompt_name} with truncated prompt: {truncated_prompt}")
+        print(f"Generating image with truncated prompt: {truncated_prompt}")
         
         # Call the model
         output = model(prompt=truncated_prompt, num_inference_steps=1, guidance_scale=0.0)
         
-        # Debugging: Print full model output
-        print(f"Full model output for {prompt_name}: {output}")
-
         # Check if output is valid
         if output is not None and hasattr(output, 'images') and output.images:
-            print(f"Image generated for {prompt_name}")
+            print(f"Image generated")
             image = output.images[0]
-            buffered = BytesIO()
-            image.save(buffered, format="JPEG")
-            image_bytes = buffered.getvalue()
-            return image_bytes
+            return image, None
         else:
-            print(f"No images found or generated output is None for {prompt_name}")
-            return None
+            print(f"No images found or generated output is None")
+            return None, "No images found or generated output is None"
     except Exception as e:
-        print(f"An error occurred while generating image for {prompt_name}: {e}")
-        return None
+        print(f"An error occurred while generating image: {e}")
+        return None, str(e)
 
-async def queue_api_calls(sentence_mapping, character_dict, selected_style):
-    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}")
+def inference(prompt):
+    print(f"Received prompt: {prompt}")  # Debugging statement
+    image, error = generate_image(prompt)
+    if error:
+        print(f"Error generating image: {error}")  # Debugging statement
+        return "Error: " + error
 
-    loop = asyncio.get_running_loop()
-    with ProcessPoolExecutor() as pool:
-        tasks = [
-            loop.run_in_executor(pool, generate_image, prompt, f"Prompt {paragraph_number}")
-            for paragraph_number, prompt in prompts
-        ]
-        responses = await asyncio.gather(*tasks)
-    
-    images = {paragraph_number: response for (paragraph_number, _), response in zip(prompts, responses)}
-    print("Finished queuing API calls. Generated images: ", images)
-    return images
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
 
 def process_prompt(sentence_mapping, character_dict, selected_style):
     print("Processing prompt...")
@@ -70,14 +54,20 @@ def process_prompt(sentence_mapping, character_dict, selected_style):
     print(f"Character Dict: {character_dict}")
     print(f"Selected Style: {selected_style}")
 
-    # Ensure we are in the right event loop context
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    cmpt_return = loop.run_until_complete(queue_api_calls(sentence_mapping, character_dict, selected_style))
-    loop.close()
+    prompts = []
+    for paragraph_number, sentences in sentence_mapping.items():
+        combined_sentence = " ".join(sentences)
+        prompt, negative_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}")
+
+    images = {}
+    for paragraph_number, prompt in prompts:
+        img_str = inference(prompt)
+        images[paragraph_number] = img_str
 
-    print("Prompt processing complete. Generated images: ", cmpt_return)
-    return cmpt_return
+    print("Prompt processing complete. Generated images: ", images)
+    return images
 
 gradio_interface = gr.Interface(
     fn=process_prompt,