Update app.py

app.py

@@ -4,13 +4,25 @@ from PIL import Image, ImageDraw, ImageFont
 import textwrap
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
+import random
+import time
 
-def download_image(prompt, index):
+# Constants for image styles
+IMAGE_STYLES = {
+    "Realistic": "",
+    "Cartoon": "cartoon style",
+    "Watercolor": "watercolor painting style",
+    "Sketch": "pencil sketch style",
+    "Anime": "anime style"
+}
+
+def download_image(prompt, index, style):
     try:
-        url = f"https://pollinations.ai/p/{prompt}"
-        response = requests.get(url, timeout=10)
+        style_prompt = f"{prompt}, {IMAGE_STYLES[style]}" if style != "Realistic" else prompt
+        url = f"https://pollinations.ai/p/{style_prompt.replace(' ', '_')}"
+        response = requests.get(url, timeout=15)
         response.raise_for_status()
-        image_filename = f'generated_image_{index}.jpg'
+        image_filename = f'generated_image_{index}_{int(time.time())}.jpg'
         with open(image_filename, 'wb') as file:
             file.write(response.content)
         return image_filename
@@ -18,45 +30,38 @@ def download_image(prompt, index):
         print(f"Error downloading image {index}: {e}")
         return None
 
-def add_text_to_image(image_path, text):
+def add_text_to_image(image_path, text, font_size, text_color):
     try:
         img = Image.open(image_path)
         draw = ImageDraw.Draw(img)
         
-        # Use a default font if custom font is not available
         try:
-            font = ImageFont.truetype("arial.ttf", 30)
+            font = ImageFont.truetype("arial.ttf", font_size)
         except IOError:
-            font = ImageFont.load_default().font_variant(size=30)
+            font = ImageFont.load_default().font_variant(size=font_size)
         
-        # Calculate text size and position
         img_width, img_height = img.size
-        max_text_width = int(img_width * 0.9)  # 90% of image width
+        max_text_width = int(img_width * 0.9)
         lines = textwrap.wrap(text, width=40)
         
-        # Use getbbox instead of getsize
-        line_height = font.getbbox("hg")[3] + 5  # Add some padding
+        line_height = font.getbbox("hg")[3] + 5
         text_height = line_height * len(lines)
         
-        # Create semi-transparent background
         overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
         overlay_draw = ImageDraw.Draw(overlay)
         overlay_draw.rectangle([(0, img_height - text_height - 20), (img_width, img_height)], 
                                fill=(0, 0, 0, 180))
         
-        # Paste the overlay onto the original image
         img = img.convert('RGBA')
         img = Image.alpha_composite(img, overlay)
         draw = ImageDraw.Draw(img)
         
-        # Add text
         y_text = img_height - text_height - 10
         for line in lines:
-            # Use getbbox for line width calculation
             bbox = font.getbbox(line)
             line_width = bbox[2] - bbox[0]
             x_text = (img_width - line_width) / 2
-            draw.text((x_text, y_text), line, font=font, fill=(255, 255, 255, 255))
+            draw.text((x_text, y_text), line, font=font, fill=text_color)
             y_text += line_height
         
         img_with_text = f'text_added_{os.path.basename(image_path)}'
@@ -67,25 +72,15 @@ def add_text_to_image(image_path, text):
         print(f"Error adding text to image: {e}")
         return None
 
-    with ThreadPoolExecutor(max_workers=5) as executor:
-        future_to_idx = {executor.submit(process_line, idx, line): idx for idx, line in enumerate(lines)}
-        for future in progress.tqdm(as_completed(future_to_idx), total=len(lines), desc="Processing images"):
-            idx = future_to_idx[future]
-            result = future.result()
-            if result:
-                images_with_text.append((idx, result))
-    
-    return [img for _, img in sorted(images_with_text)]
-
-def process_story_lines(story, progress=gr.Progress()):
+def process_story_lines(story, style, font_size, text_color, progress=gr.Progress()):
     lines = [line.strip() for line in story.split('\n') if line.strip()]
     images_with_text = []
     
     def process_line(idx, line):
         prompt = line.replace(" ", "_")
-        img_file = download_image(prompt, idx)
+        img_file = download_image(prompt, idx, style)
         if img_file:
-            return add_text_to_image(img_file, line)
+            return add_text_to_image(img_file, line, font_size, text_color)
         return None
 
     with ThreadPoolExecutor(max_workers=5) as executor:
@@ -98,24 +93,60 @@ def process_story_lines(story, progress=gr.Progress()):
     
     return [img for _, img in sorted(images_with_text)]
 
-def visualize_story_images(story):
+def visualize_story_images(story, style, font_size, text_color):
     if not story.strip():
         return []
-    return process_story_lines(story)
+    return process_story_lines(story, style, font_size, text_color)
+
+def preview_image(story, style, font_size, text_color):
+    if not story.strip():
+        return None
+    lines = [line.strip() for line in story.split('\n') if line.strip()]
+    if not lines:
+        return None
+    preview_line = random.choice(lines)
+    img_file = download_image(preview_line, 0, style)
+    if img_file:
+        return add_text_to_image(img_file, preview_line, font_size, text_color)
+    return None
 
-with gr.Blocks() as iface:
-    gr.Markdown("# Story Visualizer")
+with gr.Blocks(css="#gallery {height: 600px}") as iface:
+    gr.Markdown("# 📚✨ Story Visualizer")
     
     with gr.Row():
-        story_input = gr.Textbox(lines=10, placeholder="Enter your story here...", label="Story Input")
-        visualize_button = gr.Button("Visualize Story")
+        with gr.Column(scale=2):
+            story_input = gr.Textbox(lines=10, placeholder="Enter your story here...", label="Story Input")
+        with gr.Column(scale=1):
+            style_dropdown = gr.Dropdown(choices=list(IMAGE_STYLES.keys()), value="Realistic", label="Image Style")
+            font_size_slider = gr.Slider(minimum=20, maximum=60, value=30, step=2, label="Font Size")
+            text_color_picker = gr.ColorPicker(value="#FFFFFF", label="Text Color")
+            preview_button = gr.Button("Preview", variant="secondary")
+            visualize_button = gr.Button("Visualize Story", variant="primary")
     
+    preview_output = gr.Image(label="Preview", height=300)
     gallery_output = gr.Gallery(label="Visualized Story", show_label=False, elem_id="gallery", columns=2, rows=2, height="auto")
     
+    preview_button.click(
+        fn=preview_image,
+        inputs=[story_input, style_dropdown, font_size_slider, text_color_picker],
+        outputs=[preview_output]
+    )
+    
     visualize_button.click(
         fn=visualize_story_images,
-        inputs=[story_input],
+        inputs=[story_input, style_dropdown, font_size_slider, text_color_picker],
         outputs=[gallery_output]
     )
 
+    gr.Markdown("""
+    ## How to use:
+    1. Enter your story in the text box, with each line representing a scene.
+    2. Choose an image style from the dropdown menu.
+    3. Adjust the font size and text color as desired.
+    4. Click "Preview" to see how a random scene from your story might look.
+    5. When you're happy with the settings, click "Visualize Story" to generate images for all scenes.
+    
+    Enjoy bringing your story to life! 🎨📖
+    """)
+
 iface.launch(share=True)