testing larger layouts

app.py

@@ -68,6 +68,7 @@ def get_layout_choices(num_images: int) -> List[Tuple[str, str]]:
     key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
     if key in PAGE_LAYOUTS:
         return [(layout["label"], layout["id"]) for layout in PAGE_LAYOUTS[key]]
+    # Return empty list if no layouts found (shouldn't happen with our config)
     return [("Default", "default")]
 
 def get_random_style_preset():
@@ -252,25 +253,96 @@ pipe.fuse_lora()
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 
-def get_image_size(aspect_ratio):
-    """Converts aspect ratio string to width, height tuple, optimized for 1024 base."""
-    if aspect_ratio == "1:1":
-        return 1024, 1024
-    elif aspect_ratio == "16:9":
-        return 1152, 640
-    elif aspect_ratio == "9:16":
-        return 640, 1152
-    elif aspect_ratio == "4:3":
-        return 1024, 768
-    elif aspect_ratio == "3:4":
-        return 768, 1024
-    elif aspect_ratio == "3:2":
-        return 1024, 688
-    elif aspect_ratio == "2:3":
-        return 688, 1024
-    else:
-        # Default to 1:1 if something goes wrong
-        return 1024, 1024
+def get_image_size_for_position(position_data, image_index, num_images):
+    """Determines optimal image size based on its position in the layout.
+
+    Args:
+        position_data: Layout position data [x, y, width, height] in relative units
+        image_index: Index of the current image (0-based)
+        num_images: Total number of images in the layout
+
+    Returns:
+        tuple: (width, height) optimized for the position's aspect ratio, max 1024 in any dimension
+    """
+    if not position_data:
+        return 1024, 1024  # Default square
+
+    x_rel, y_rel, w_rel, h_rel = position_data
+    aspect_ratio = w_rel / h_rel if h_rel > 0 else 1.0
+
+    # Max dimension is 1024
+    max_dim = 1024
+
+    # Calculate dimensions maintaining aspect ratio with max of 1024
+    if aspect_ratio >= 1:  # Wider than tall
+        width = max_dim
+        height = int(max_dim / aspect_ratio)
+        # Ensure height is at least 256 for quality
+        if height < 256:
+            height = 256
+            width = int(256 * aspect_ratio)
+    else:  # Taller than wide
+        height = max_dim
+        width = int(max_dim * aspect_ratio)
+        # Ensure width is at least 256 for quality
+        if width < 256:
+            width = 256
+            height = int(256 / aspect_ratio)
+
+    # Round to nearest 64 for better compatibility
+    width = (width // 64) * 64
+    height = (height // 64) * 64
+
+    # Ensure we don't exceed max_dim after rounding
+    if width > max_dim:
+        width = max_dim
+    if height > max_dim:
+        height = max_dim
+
+    # Minimum size check
+    width = max(width, 256)
+    height = max(height, 256)
+
+    return width, height
+
+def get_layout_position_for_image(layout_id, num_images, image_index):
+    """Get the position data for a specific image in a layout.
+
+    Args:
+        layout_id: ID of the selected layout
+        num_images: Total number of images
+        image_index: Index of the current image (0-based)
+
+    Returns:
+        Position data [x, y, width, height] or None
+    """
+    key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
+    layouts = PAGE_LAYOUTS.get(key, [])
+    layout = next((l for l in layouts if l["id"] == layout_id), None)
+
+    if layout and "positions" in layout:
+        positions = layout["positions"]
+        if image_index < len(positions):
+            return positions[image_index]
+
+    # Fallback positions for each number of images
+    fallback_positions = {
+        1: [[0.05, 0.05, 0.9, 0.9]],
+        2: [[0.05, 0.05, 0.425, 0.9], [0.525, 0.05, 0.425, 0.9]],
+        3: [[0.05, 0.25, 0.283, 0.5], [0.358, 0.25, 0.283, 0.5], [0.666, 0.25, 0.283, 0.5]],
+        4: [[0.05, 0.05, 0.425, 0.425], [0.525, 0.05, 0.425, 0.425],
+            [0.05, 0.525, 0.425, 0.425], [0.525, 0.525, 0.425, 0.425]],
+        5: [[0.05, 0.05, 0.9, 0.3], [0.05, 0.4, 0.283, 0.55], [0.358, 0.4, 0.283, 0.55],
+            [0.666, 0.4, 0.283, 0.275], [0.666, 0.7, 0.283, 0.275]],
+        6: [[0.05, 0.05, 0.425, 0.283], [0.525, 0.05, 0.425, 0.283],
+            [0.05, 0.358, 0.425, 0.283], [0.525, 0.358, 0.425, 0.283],
+            [0.05, 0.666, 0.425, 0.283], [0.525, 0.666, 0.425, 0.283]]
+    }
+
+    positions = fallback_positions.get(num_images, fallback_positions[1])
+    if image_index < len(positions):
+        return positions[image_index]
+    return [0.05, 0.05, 0.9, 0.9]  # Ultimate default
 
 # --- Session Management Functions ---
 
@@ -398,9 +470,19 @@ def create_single_page_pdf(images: List[Image.Image], layout_id: str, num_images
             positions = [[0.05, 0.05, 0.425, 0.9], [0.525, 0.05, 0.425, 0.9]]
         elif num_images == 3:
             positions = [[0.05, 0.05, 0.283, 0.9], [0.358, 0.05, 0.283, 0.9], [0.666, 0.05, 0.283, 0.9]]
-        else:
+        elif num_images == 4:
             positions = [[0.05, 0.05, 0.425, 0.425], [0.525, 0.05, 0.425, 0.425],
                         [0.05, 0.525, 0.425, 0.425], [0.525, 0.525, 0.425, 0.425]]
+        elif num_images == 5:
+            positions = [[0.05, 0.05, 0.9, 0.3], [0.05, 0.4, 0.283, 0.55], [0.358, 0.4, 0.283, 0.55],
+                        [0.666, 0.4, 0.283, 0.275], [0.666, 0.7, 0.283, 0.275]]
+        elif num_images == 6:
+            positions = [[0.05, 0.05, 0.425, 0.283], [0.525, 0.05, 0.425, 0.283],
+                        [0.05, 0.358, 0.425, 0.283], [0.525, 0.358, 0.425, 0.283],
+                        [0.05, 0.666, 0.425, 0.283], [0.525, 0.666, 0.425, 0.283]]
+        else:
+            # For more than 6, create a simple grid
+            positions = [[0.05, 0.05, 0.9, 0.9]]
     else:
         positions = layout["positions"]
 
@@ -424,8 +506,9 @@ def create_single_page_pdf(images: List[Image.Image], layout_id: str, num_images
         image.save(img_buffer, format='JPEG', quality=95)
         img_buffer.seek(0)
 
-        # Draw the image on the PDF
-        pdf.drawImage(ImageReader(img_buffer), x, y, width=width, height=height, preserveAspectRatio=True)
+        # Draw the image on the PDF - fill the entire allocated space
+        # preserveAspectRatio=False to fill the entire area
+        pdf.drawImage(ImageReader(img_buffer), x, y, width=width, height=height, preserveAspectRatio=False, mask='auto')
 
     # Save the PDF
     pdf.save()
@@ -473,12 +556,11 @@ def create_multi_page_pdf(session_manager: SessionManager) -> str:
     return str(pdf_path)
 
 # --- Main Inference Function (with session support) ---
-@spaces.GPU(duration=120)  # Increased duration for multiple images
+@spaces.GPU(duration=180)  # Increased duration for up to 6 images
 def infer_page(
     prompt,
     seed=42,
     randomize_seed=False,
-    aspect_ratio="1:1",
     guidance_scale=1.0,
     num_inference_steps=8,
     prompt_enhance=True,
@@ -496,7 +578,6 @@ def infer_page(
         prompt (str): The text prompt to generate images from.
         seed (int): The seed for the random number generator for reproducibility.
         randomize_seed (bool): If True, a random seed is used for each image.
-        aspect_ratio (str): The desired aspect ratio of the output images.
         guidance_scale (float): Corresponds to `true_cfg_scale`.
         num_inference_steps (int): The number of denoising steps.
         prompt_enhance (bool): If True, the prompt is rewritten by an external LLM.
@@ -533,12 +614,17 @@ def infer_page(
 
         current_seed = seed + i if not randomize_seed else random.randint(0, MAX_SEED)
 
-        # Generate single image
-        image, used_seed = infer_single(
+        # Get optimal aspect ratio based on position in layout
+        position_data = get_layout_position_for_image(layout, int(num_images), i)
+
+        # Generate single image with automatic aspect ratio
+        image, used_seed = infer_single_auto(
             prompt=prompt,
             seed=current_seed,
             randomize_seed=False,  # We handle randomization here
-            aspect_ratio=aspect_ratio,
+            position_data=position_data,
+            image_index=i,
+            num_images=int(num_images),
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             prompt_enhance=prompt_enhance,
@@ -572,12 +658,14 @@ def infer_page(
 
     return session_state, pdf_path, generated_images[0] if generated_images else None, page_info, button_label
 
-# Rename the original infer function
-def infer_single(
+# New inference function with automatic aspect ratio
+def infer_single_auto(
     prompt,
     seed=42,
     randomize_seed=False,
-    aspect_ratio="1:1",
+    position_data=None,
+    image_index=0,
+    num_images=1,
     guidance_scale=1.0,
     num_inference_steps=8,
     prompt_enhance=True,
@@ -585,39 +673,20 @@ def infer_single(
     custom_style_text="",
 ):
     """
-    Generates an image based on a text prompt using the Qwen-Image-Lightning model.
-
-    Args:
-        prompt (str): The text prompt to generate the image from.
-        seed (int): The seed for the random number generator for reproducibility.
-        randomize_seed (bool): If True, a random seed is used.
-        aspect_ratio (str): The desired aspect ratio of the output image.
-        guidance_scale (float): Corresponds to `true_cfg_scale`. A higher value
-            encourages the model to generate images that are more closely related
-            to the prompt.
-        num_inference_steps (int): The number of denoising steps.
-        prompt_enhance (bool): If True, the prompt is rewritten by an external
-            LLM to add more detail.
-        style_preset (str): The key of the style preset to apply.
-        custom_style_text (str): Custom style text when 'no_style' is selected.
-        progress (gr.Progress): A Gradio Progress object to track the generation
-            progress in the UI.
-
-    Returns:
-        tuple[Image.Image, int]: A tuple containing the generated PIL Image and
-            the integer seed used for the generation.
+    Generates an image with automatically determined aspect ratio based on layout position.
     """
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Convert aspect ratio to width and height
-    width, height = get_image_size(aspect_ratio)
+    # Automatically determine image size based on position
+    width, height = get_image_size_for_position(position_data, image_index, num_images)
 
     # Set up the generator for reproducibility
     generator = torch.Generator(device="cuda").manual_seed(seed)
 
     print(f"Original prompt: '{prompt}'")
     print(f"Style preset: '{style_preset}'")
+    print(f"Auto-selected size based on layout: {width}x{height}")
 
     # Apply style preset first
     styled_prompt, style_negative_prompt = apply_style_preset(prompt, style_preset, custom_style_text)
@@ -646,8 +715,8 @@ def infer_single(
 
     return image, seed
 
-# Keep the old infer function for backward compatibility
-infer = infer_single
+# Keep the old infer function for backward compatibility (simplified)
+infer = infer_single_auto
 
 # --- Examples and UI Layout ---
 examples = [
@@ -698,6 +767,46 @@ with gr.Blocks(css=css) as demo:
                 run_button = gr.Button("Generate page 1", variant="primary")
                 reset_button = gr.Button("Start New Document", variant="secondary")
 
+        # New row for Style Preset and Page Layout
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Number of images slider (affects layout choices)
+                num_images_slider = gr.Slider(
+                    label="Images per page",
+                    minimum=1,
+                    maximum=6,
+                    step=1,
+                    value=1,
+                    info="Number of images to generate for the PDF (1-6)"
+                )
+
+            with gr.Column(scale=2):
+                layout_dropdown = gr.Dropdown(
+                    label="Page Layout",
+                    choices=[("Full Page", "full_page")],
+                    value="full_page",
+                    interactive=True,
+                    info="How images are arranged on the page"
+                )
+
+            with gr.Column(scale=2):
+                # Create dropdown choices from loaded presets
+                style_choices = [(preset["label"], key) for key, preset in STYLE_PRESETS.items()]
+                style_preset = gr.Dropdown(
+                    label="Style Preset",
+                    choices=style_choices,
+                    value="no_style",
+                    interactive=True
+                )
+
+            with gr.Column(scale=2):
+                custom_style_text = gr.Textbox(
+                    label="Custom Style Text",
+                    placeholder="Enter custom style (e.g., 'oil painting')",
+                    visible=False,
+                    lines=1
+                )
+
         with gr.Row():
             with gr.Column(scale=1):
                 result_preview = gr.Image(label="Preview", show_label=True, type="pil")
@@ -719,48 +828,8 @@ with gr.Blocks(css=css) as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
             with gr.Row():
-                aspect_ratio = gr.Radio(
-                    label="Aspect ratio (width:height)",
-                    choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
-                    value="16:9",
-                )
                 prompt_enhance = gr.Checkbox(label="Prompt Enhance", value=True)
 
-            with gr.Row():
-                # Create dropdown choices from loaded presets
-                style_choices = [(preset["label"], key) for key, preset in STYLE_PRESETS.items()]
-                style_preset = gr.Dropdown(
-                    label="Style Preset",
-                    choices=style_choices,
-                    value="no_style",
-                    interactive=True
-                )
-
-                custom_style_text = gr.Textbox(
-                    label="Custom Style Text",
-                    placeholder="Enter custom style keywords (e.g., 'oil painting, impressionist')",
-                    visible=False,
-                    lines=2
-                )
-
-            with gr.Row():
-                num_images_slider = gr.Slider(
-                    label="Images per page",
-                    minimum=1,
-                    maximum=4,
-                    step=1,
-                    value=1,
-                    info="Number of images to generate for the PDF"
-                )
-
-                layout_dropdown = gr.Dropdown(
-                    label="Page Layout",
-                    choices=[("Full Page", "full_page")],
-                    value="full_page",
-                    interactive=True,
-                    info="How images are arranged on the page"
-                )
-
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale (True CFG Scale)",
@@ -778,33 +847,35 @@ with gr.Blocks(css=css) as demo:
                     value=8,
                 )
 
-            # Add interaction to show/hide custom style text field
-            def toggle_custom_style(style_value):
-                return gr.update(visible=(style_value == "no_style"))
+        # Add interaction to show/hide custom style text field
+        def toggle_custom_style(style_value):
+            return gr.update(visible=(style_value == "no_style"))
 
-            style_preset.change(
-                fn=toggle_custom_style,
-                inputs=[style_preset],
-                outputs=[custom_style_text]
-            )
+        style_preset.change(
+            fn=toggle_custom_style,
+            inputs=[style_preset],
+            outputs=[custom_style_text]
+        )
 
-            # Update layout dropdown when number of images changes
-            def update_layout_choices(num_images):
-                choices = get_layout_choices(int(num_images))
-                return gr.update(choices=choices, value=choices[0][1] if choices else "default")
+        # Update layout dropdown when number of images changes
+        def update_layout_choices(num_images):
+            choices = get_layout_choices(int(num_images))
+            return gr.update(choices=choices, value=choices[0][1] if choices else "default")
 
-            num_images_slider.change(
-                fn=update_layout_choices,
-                inputs=[num_images_slider],
-                outputs=[layout_dropdown]
-            )
+        num_images_slider.change(
+            fn=update_layout_choices,
+            inputs=[num_images_slider],
+            outputs=[layout_dropdown]
+        )
 
         # Update examples to show some with different styles and image counts
         styled_examples = [
             ["A capybara wearing a suit holding a sign that reads Hello World", "no_style", "", 1],
-            ["sharks raining down on san francisco", "flying_saucer", "", 2],
-            ["A beautiful landscape with mountains and a lake", "klimt", "", 3],
+            ["sharks raining down on san francisco", "anime", "", 2],
+            ["A beautiful landscape with mountains and a lake", "watercolor", "", 3],
             ["A knight fighting a dragon", "medieval", "", 4],
+            ["Space battle with laser beams", "sci-fi", "", 5],
+            ["Detective investigating a mystery", "noir", "", 6],
         ]
 
         gr.Examples(
@@ -823,7 +894,6 @@ with gr.Blocks(css=css) as demo:
             prompt,
             seed,
             randomize_seed,
-            aspect_ratio,
             guidance_scale,
             num_inference_steps,
             prompt_enhance,

page_layouts.yaml

@@ -124,4 +124,122 @@ layouts:
         - [0.05, 0.05, 0.425, 0.425]  # Top left (large)
         - [0.525, 0.05, 0.425, 0.425]  # Top right (large)
         - [0.05, 0.525, 0.425, 0.425]  # Bottom left
-        - [0.525, 0.7, 0.425, 0.25]  # Bottom right (small)
+        - [0.525, 0.7, 0.425, 0.25]  # Bottom right (small)
+
+  5_images:
+    - id: "us_comic_action"
+      label: "US Comic - Action Scene"
+      description: "Classic American superhero comic layout with large establishing shot"
+      positions:
+        - [0.05, 0.05, 0.9, 0.4]     # Wide establishing shot (panoramic)
+        - [0.05, 0.5, 0.283, 0.45]   # Action panel 1
+        - [0.358, 0.5, 0.283, 0.45]  # Action panel 2
+        - [0.666, 0.5, 0.283, 0.225] # Close-up 1
+        - [0.666, 0.75, 0.283, 0.225] # Close-up 2
+
+    - id: "manga_vertical_flow"
+      label: "Manga - Vertical Flow"
+      description: "Japanese manga style with vertical reading flow"
+      positions:
+        - [0.525, 0.05, 0.425, 0.35]  # Top right (read first in manga)
+        - [0.05, 0.05, 0.425, 0.35]   # Top left
+        - [0.525, 0.45, 0.425, 0.25]  # Middle right
+        - [0.05, 0.45, 0.425, 0.25]   # Middle left
+        - [0.05, 0.75, 0.9, 0.2]      # Bottom wide panel
+
+    - id: "euro_bd_grid"
+      label: "European BD - Clear Grid"
+      description: "Franco-Belgian clear line style with regular panels"
+      positions:
+        - [0.05, 0.05, 0.425, 0.283]  # Row 1 left
+        - [0.525, 0.05, 0.425, 0.283] # Row 1 right
+        - [0.05, 0.358, 0.9, 0.283]   # Row 2 wide
+        - [0.05, 0.666, 0.425, 0.283] # Row 3 left
+        - [0.525, 0.666, 0.425, 0.283] # Row 3 right
+
+    - id: "diagonal_dynamic"
+      label: "Dynamic Diagonal"
+      description: "Action-oriented diagonal composition"
+      positions:
+        - [0.05, 0.05, 0.5, 0.4]    # Large top left
+        - [0.6, 0.05, 0.35, 0.25]   # Small top right
+        - [0.3, 0.35, 0.4, 0.3]     # Center focus
+        - [0.05, 0.7, 0.35, 0.25]   # Bottom left
+        - [0.6, 0.7, 0.35, 0.25]    # Bottom right
+
+    - id: "spiral_focus"
+      label: "Spiral Focus"
+      description: "Panels arranged in a spiral leading to center"
+      positions:
+        - [0.05, 0.05, 0.35, 0.35]   # Top left
+        - [0.425, 0.05, 0.525, 0.25] # Top wide
+        - [0.7, 0.35, 0.25, 0.6]     # Right tall
+        - [0.425, 0.7, 0.525, 0.25]  # Bottom wide
+        - [0.25, 0.35, 0.4, 0.3]     # Center focus
+
+  6_images:
+    - id: "classic_comic_grid"
+      label: "Classic Comic Grid"
+      description: "Traditional 2x3 American comic book grid"
+      positions:
+        - [0.05, 0.05, 0.425, 0.283]   # Row 1 left
+        - [0.525, 0.05, 0.425, 0.283]  # Row 1 right
+        - [0.05, 0.358, 0.425, 0.283]  # Row 2 left
+        - [0.525, 0.358, 0.425, 0.283] # Row 2 right
+        - [0.05, 0.666, 0.425, 0.283]  # Row 3 left
+        - [0.525, 0.666, 0.425, 0.283] # Row 3 right
+
+    - id: "manga_4koma"
+      label: "Manga - 4-Koma Plus"
+      description: "Japanese 4-panel strip with header and footer"
+      positions:
+        - [0.05, 0.05, 0.9, 0.15]    # Header panel
+        - [0.05, 0.25, 0.425, 0.2]   # Strip 1
+        - [0.525, 0.25, 0.425, 0.2]  # Strip 2
+        - [0.05, 0.5, 0.425, 0.2]    # Strip 3
+        - [0.525, 0.5, 0.425, 0.2]   # Strip 4
+        - [0.05, 0.75, 0.9, 0.2]     # Footer/punchline
+
+    - id: "euro_bd_cinematic"
+      label: "European BD - Cinematic"
+      description: "Cinematic European style with varied panel sizes"
+      positions:
+        - [0.05, 0.05, 0.9, 0.25]    # Wide establishing
+        - [0.05, 0.35, 0.283, 0.25]  # Small 1
+        - [0.358, 0.35, 0.283, 0.25] # Small 2
+        - [0.666, 0.35, 0.283, 0.25] # Small 3
+        - [0.05, 0.65, 0.425, 0.3]   # Medium left
+        - [0.525, 0.65, 0.425, 0.3]  # Medium right
+
+    - id: "action_sequence"
+      label: "Action Sequence"
+      description: "Fast-paced action scene layout"
+      positions:
+        - [0.05, 0.05, 0.6, 0.35]    # Large action shot
+        - [0.7, 0.05, 0.25, 0.175]   # Speed line 1
+        - [0.7, 0.25, 0.25, 0.175]   # Speed line 2
+        - [0.05, 0.45, 0.283, 0.5]   # Vertical impact 1
+        - [0.358, 0.45, 0.283, 0.5]  # Vertical impact 2
+        - [0.666, 0.45, 0.283, 0.5]  # Vertical impact 3
+
+    - id: "storytelling_flow"
+      label: "Storytelling Flow"
+      description: "Natural reading flow for narrative scenes"
+      positions:
+        - [0.05, 0.05, 0.425, 0.25]  # Scene 1
+        - [0.525, 0.05, 0.425, 0.25] # Scene 2
+        - [0.05, 0.35, 0.9, 0.2]     # Wide transition
+        - [0.05, 0.6, 0.425, 0.35]   # Scene 3
+        - [0.525, 0.6, 0.425, 0.175] # Scene 4a
+        - [0.525, 0.8, 0.425, 0.175] # Scene 4b
+
+    - id: "focus_surround"
+      label: "Focus with Details"
+      description: "Central focus with surrounding detail panels"
+      positions:
+        - [0.25, 0.25, 0.5, 0.5]     # Large center focus
+        - [0.05, 0.05, 0.35, 0.15]   # Top left detail
+        - [0.6, 0.05, 0.35, 0.15]    # Top right detail
+        - [0.05, 0.8, 0.35, 0.15]    # Bottom left detail
+        - [0.6, 0.8, 0.35, 0.15]     # Bottom right detail
+        - [0.05, 0.4, 0.15, 0.3]     # Left side detail