rethinking this project to be Gradio+API+PDF based

CLAUDE.md

@@ -0,0 +1,92 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This is a Gradio application for image generation using the Qwen-Image model with Lightning LoRA acceleration. It's designed to run on Hugging Face Spaces with GPU support, providing fast 8-step image generation with advanced text rendering capabilities.
+
+## Commands
+
+### Run the application locally
+```bash
+python app.py
+```
+
+### Install dependencies
+```bash
+pip install -r requirements.txt
+```
+
+## Architecture
+
+### Core Components
+
+1. **Model Pipeline** (`app.py:130-164`)
+   - Uses `Qwen/Qwen-Image` diffusion model with custom FlowMatchEulerDiscreteScheduler
+   - Loads Lightning LoRA weights for 8-step acceleration
+   - Configured for bfloat16 precision on CUDA
+
+2. **Prompt Enhancement System** (`app.py:41-125`)
+   - `polish_prompt()`: Uses Hugging Face InferenceClient with Cerebras provider to enhance prompts
+   - `get_caption_language()`: Detects Chinese vs English prompts
+   - `rewrite()`: Language-specific prompt enhancement with different system prompts for Chinese/English
+   - Requires `HF_TOKEN` environment variable for API access
+
+3. **Style Presets System** (`app.py:16-87`)
+   - `load_style_presets()`: Loads style presets from `style_presets.yaml`
+   - `apply_style_preset()`: Applies selected style to prompts
+   - Supports custom styles and random style selection
+   - Each preset includes prefix, suffix, and negative prompt components
+
+4. **Page Layouts System** (`app.py:89-111`)
+   - `load_page_layouts()`: Loads multi-image layouts from `page_layouts.yaml`
+   - Supports 1-4 images per page with various layout configurations
+   - Dynamic layout selection based on number of images
+
+5. **PDF Generation** (`app.py:166-223`)
+   - `create_pdf_with_layout()`: Creates PDF with multiple images in selected layout
+   - Uses ReportLab for high-quality PDF generation
+   - Preserves image quality at 95% JPEG compression
+   - A4 page size with flexible positioning system
+
+6. **Multi-Image Generation** (`app.py:225-307`)
+   - `infer_multiple()`: Generates multiple images and combines into PDF
+   - Progressive generation with status updates
+   - Seed management for reproducibility across multiple images
+   - Returns PDF file, preview image, and seed information
+
+7. **Gradio Interface** (`app.py:380-500+`)
+   - Slider for selecting 1-4 images per page
+   - Dynamic layout dropdown that updates based on image count
+   - Style preset dropdown with custom style text option
+   - PDF download and image preview outputs
+   - Advanced settings for all generation parameters
+
+## Key Configuration
+
+- **Scheduler Config** (`app.py:133-148`): Custom configuration for FlowMatchEulerDiscreteScheduler with exponential time shifting
+- **Aspect Ratios** (`app.py:170-188`): Predefined aspect ratios optimized for 1024 base resolution
+- **Style Presets** (`style_presets.yaml`): Configurable style presets with prompt modifiers and negative prompts
+- **Page Layouts** (`page_layouts.yaml`): Flexible layout system for 1-4 images per page
+- **Default Settings**: 8 inference steps, guidance scale 1.0, prompt enhancement enabled, 1 image per page
+
+## Environment Variables
+
+- `HF_TOKEN`: Required for prompt enhancement via Hugging Face InferenceClient
+- Used for accessing Cerebras provider for Qwen3-235B model
+
+## Key Features
+
+- **Session-based storage**: Each user session gets a unique temporary directory that persists for 24 hours
+- **Multi-page PDF generation**: Users can generate up to 128 pages in a single document
+- **Dynamic page addition**: Click "Generate page N" to add the next page to the PDF
+- **Flexible layouts**: Different layout options for 1-4 images per page
+- **Style presets**: 20+ predefined artistic styles
+- **Automatic cleanup**: Old sessions are automatically cleaned after 24 hours
+
+## Model Dependencies
+
+- Main model: `Qwen/Qwen-Image`
+- LoRA weights: `lightx2v/Qwen-Image-Lightning` (V1.1 safetensors)
+- Prompt enhancement model: `Qwen/Qwen3-235B-A22B-Instruct-2507` via Cerebras

app.py

@@ -5,10 +5,121 @@ import torch
 import spaces
 import math
 import os
+import yaml
+import io
+import tempfile
+import shutil
+import uuid
+import time
+import json
+from typing import List, Tuple, Dict, Optional
+from datetime import datetime, timedelta
+from pathlib import Path
 
 from PIL import Image
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
 from huggingface_hub import InferenceClient
+from reportlab.lib.pagesizes import A4
+from reportlab.pdfgen import canvas
+from reportlab.pdfbase import pdfmetrics
+from reportlab.lib.utils import ImageReader
+from PyPDF2 import PdfReader, PdfWriter
+
+# --- Style Presets Loading ---
+
+def load_style_presets():
+    """Load style presets from YAML file."""
+    try:
+        with open('style_presets.yaml', 'r') as f:
+            data = yaml.safe_load(f)
+        # Filter only enabled presets
+        presets = {k: v for k, v in data['presets'].items() if v.get('enabled', True)}
+        return presets
+    except Exception as e:
+        print(f"Error loading style presets: {e}")
+        return {"no_style": {"id": "no_style", "label": "No style (custom)", "prompt_prefix": "", "prompt_suffix": "", "negative_prompt": ""}}
+
+# Load presets at startup
+STYLE_PRESETS = load_style_presets()
+
+# --- Page Layouts Loading ---
+
+def load_page_layouts():
+    """Load page layouts from YAML file."""
+    try:
+        with open('page_layouts.yaml', 'r') as f:
+            data = yaml.safe_load(f)
+        return data['layouts']
+    except Exception as e:
+        print(f"Error loading page layouts: {e}")
+        # Fallback to basic layouts
+        return {
+            1: [{"id": "full_page", "label": "Full Page", "positions": [[0.05, 0.05, 0.9, 0.9]]}],
+            2: [{"id": "horizontal_split", "label": "Horizontal Split", "positions": [[0.05, 0.05, 0.425, 0.9], [0.525, 0.05, 0.425, 0.9]]}],
+            3: [{"id": "grid", "label": "Grid", "positions": [[0.05, 0.05, 0.283, 0.5], [0.358, 0.05, 0.283, 0.5], [0.666, 0.05, 0.283, 0.5]]}],
+            4: [{"id": "grid_2x2", "label": "2x2 Grid", "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]]}]
+        }
+
+# Load layouts at startup
+PAGE_LAYOUTS = load_page_layouts()
+
+def get_layout_choices(num_images: int) -> List[Tuple[str, str]]:
+    """Get available layout choices for a given number of images."""
+    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 [("Default", "default")]
+
+def get_random_style_preset():
+    """Get a random style preset (excluding 'no_style' and 'random')."""
+    eligible_keys = [k for k in STYLE_PRESETS.keys() if k not in ['no_style', 'random']]
+    if eligible_keys:
+        return random.choice(eligible_keys)
+    return 'no_style'
+
+def apply_style_preset(prompt, style_preset_key, custom_style_text=""):
+    """
+    Apply style preset to the prompt.
+
+    Args:
+        prompt: The user's base prompt
+        style_preset_key: The key of the selected style preset
+        custom_style_text: Custom style text when 'no_style' is selected
+
+    Returns:
+        tuple: (styled_prompt, negative_prompt)
+    """
+    if style_preset_key == 'no_style':
+        # Use custom style text if provided
+        if custom_style_text and custom_style_text.strip():
+            styled_prompt = f"{custom_style_text}, {prompt}"
+        else:
+            styled_prompt = prompt
+        return styled_prompt, ""
+
+    if style_preset_key == 'random':
+        # Select a random style
+        style_preset_key = get_random_style_preset()
+
+    if style_preset_key in STYLE_PRESETS:
+        preset = STYLE_PRESETS[style_preset_key]
+        prefix = preset.get('prompt_prefix', '')
+        suffix = preset.get('prompt_suffix', '')
+        negative = preset.get('negative_prompt', '')
+
+        # Build the styled prompt
+        parts = []
+        if prefix:
+            parts.append(prefix)
+        parts.append(prompt)
+        if suffix:
+            parts.append(suffix)
+
+        styled_prompt = ', '.join(parts)
+        return styled_prompt, negative
+
+    # Fallback to original prompt if preset not found
+    return prompt, ""
 
 # --- New Prompt Enhancement using Hugging Face InferenceClient ---
 
@@ -161,9 +272,209 @@ def get_image_size(aspect_ratio):
         # Default to 1:1 if something goes wrong
         return 1024, 1024
 
-# --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU(duration=60)
-def infer(
+# --- Session Management Functions ---
+
+class SessionManager:
+    """Manages user session data and temporary file storage."""
+
+    def __init__(self, session_id: str = None):
+        self.session_id = session_id or str(uuid.uuid4())
+        self.base_dir = Path(tempfile.gettempdir()) / "gradio_comic_sessions"
+        self.session_dir = self.base_dir / self.session_id
+        self.session_dir.mkdir(parents=True, exist_ok=True)
+        self.metadata_file = self.session_dir / "metadata.json"
+        self.pdf_path = self.session_dir / "comic.pdf"
+        self.load_or_create_metadata()
+
+    def load_or_create_metadata(self):
+        """Load existing metadata or create new."""
+        if self.metadata_file.exists():
+            with open(self.metadata_file, 'r') as f:
+                self.metadata = json.load(f)
+        else:
+            self.metadata = {
+                "created_at": datetime.now().isoformat(),
+                "pages": [],
+                "total_pages": 0
+            }
+            self.save_metadata()
+
+    def save_metadata(self):
+        """Save metadata to file."""
+        with open(self.metadata_file, 'w') as f:
+            json.dump(self.metadata, f, indent=2)
+
+    def add_page(self, images: List[Image.Image], layout_id: str, seeds: List[int]):
+        """Add a new page to the session."""
+        page_num = self.metadata["total_pages"] + 1
+        page_dir = self.session_dir / f"page_{page_num}"
+        page_dir.mkdir(exist_ok=True)
+
+        # Save images
+        image_paths = []
+        for i, img in enumerate(images):
+            img_path = page_dir / f"image_{i+1}.jpg"
+            img.save(img_path, 'JPEG', quality=95)
+            image_paths.append(str(img_path))
+
+        # Update metadata
+        self.metadata["pages"].append({
+            "page_num": page_num,
+            "layout_id": layout_id,
+            "num_images": len(images),
+            "image_paths": image_paths,
+            "seeds": seeds,
+            "created_at": datetime.now().isoformat()
+        })
+        self.metadata["total_pages"] = page_num
+        self.save_metadata()
+
+        return page_num
+
+    def get_all_pages_images(self) -> List[Tuple[List[Image.Image], str, int]]:
+        """Get all images from all pages."""
+        pages_data = []
+        for page in self.metadata["pages"]:
+            images = []
+            for img_path in page["image_paths"]:
+                if Path(img_path).exists():
+                    images.append(Image.open(img_path))
+            if images:
+                pages_data.append((images, page["layout_id"], page["num_images"]))
+        return pages_data
+
+    def cleanup_old_sessions(self, max_age_hours: int = 24):
+        """Clean up sessions older than max_age_hours."""
+        if not self.base_dir.exists():
+            return
+
+        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
+
+        for session_dir in self.base_dir.iterdir():
+            if session_dir.is_dir():
+                metadata_file = session_dir / "metadata.json"
+                if metadata_file.exists():
+                    try:
+                        with open(metadata_file, 'r') as f:
+                            metadata = json.load(f)
+                        created_at = datetime.fromisoformat(metadata["created_at"])
+                        if created_at < cutoff_time:
+                            shutil.rmtree(session_dir)
+                            print(f"Cleaned up old session: {session_dir.name}")
+                    except Exception as e:
+                        print(f"Error cleaning session {session_dir.name}: {e}")
+
+# --- PDF Generation Functions ---
+
+def create_single_page_pdf(images: List[Image.Image], layout_id: str, num_images: int) -> bytes:
+    """
+    Create a single PDF page with images arranged according to the selected layout.
+
+    Args:
+        images: List of PIL images
+        layout_id: ID of the selected layout
+        num_images: Number of images to include
+
+    Returns:
+        PDF page as bytes
+    """
+    # Create a bytes buffer for the PDF
+    pdf_buffer = io.BytesIO()
+
+    # Create canvas with A4 size
+    pdf = canvas.Canvas(pdf_buffer, pagesize=A4)
+    page_width, page_height = A4
+
+    # Get the layout configuration
+    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 not layout:
+        # Fallback to default grid layout
+        if num_images == 1:
+            positions = [[0.05, 0.05, 0.9, 0.9]]
+        elif num_images == 2:
+            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:
+            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]]
+    else:
+        positions = layout["positions"]
+
+    # Draw each image according to the layout
+    for i, (image, pos) in enumerate(zip(images[:num_images], positions)):
+        if i >= len(images):
+            break
+
+        x_rel, y_rel, w_rel, h_rel = pos
+
+        # Convert relative positions to absolute positions
+        # Note: In ReportLab, y=0 is at the bottom
+        x = x_rel * page_width
+        y = (1 - y_rel - h_rel) * page_height  # Flip Y coordinate
+        width = w_rel * page_width
+        height = h_rel * page_height
+
+        # Convert PIL image to format suitable for ReportLab
+        img_buffer = io.BytesIO()
+        # Save with good quality
+        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)
+
+    # Save the PDF
+    pdf.save()
+
+    # Get the PDF bytes
+    pdf_buffer.seek(0)
+    pdf_bytes = pdf_buffer.read()
+
+    return pdf_bytes
+
+def create_multi_page_pdf(session_manager: SessionManager) -> str:
+    """
+    Create a multi-page PDF from all pages in the session.
+
+    Args:
+        session_manager: SessionManager instance with page data
+
+    Returns:
+        Path to the created PDF file
+    """
+    pages_data = session_manager.get_all_pages_images()
+
+    if not pages_data:
+        return None
+
+    # Create PDF writer
+    pdf_writer = PdfWriter()
+
+    # Create each page
+    for images, layout_id, num_images in pages_data:
+        page_pdf_bytes = create_single_page_pdf(images, layout_id, num_images)
+
+        # Read the single page PDF
+        page_pdf_reader = PdfReader(io.BytesIO(page_pdf_bytes))
+
+        # Add the page to the writer
+        for page in page_pdf_reader.pages:
+            pdf_writer.add_page(page)
+
+    # Write to file
+    pdf_path = session_manager.pdf_path
+    with open(pdf_path, 'wb') as f:
+        pdf_writer.write(f)
+
+    return str(pdf_path)
+
+# --- Main Inference Function (with session support) ---
+@spaces.GPU(duration=120)  # Increased duration for multiple images
+def infer_page(
     prompt,
     seed=42,
     randomize_seed=False,
@@ -171,7 +482,107 @@ def infer(
     guidance_scale=1.0,
     num_inference_steps=8,
     prompt_enhance=True,
+    style_preset="no_style",
+    custom_style_text="",
+    num_images=1,
+    layout="default",
+    session_state=None,
     progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Generates images for a new page and adds them to the PDF.
+
+    Args:
+        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.
+        style_preset (str): The key of the style preset to apply.
+        custom_style_text (str): Custom style text when 'no_style' is selected.
+        num_images (int): Number of images to generate (1-4).
+        layout (str): The layout ID for arranging images in the PDF.
+        session_state: Current session state dictionary.
+        progress (gr.Progress): A Gradio Progress object to track generation.
+
+    Returns:
+        tuple: Updated session state, PDF path, preview image, page info, and updated button label.
+    """
+    # Initialize or retrieve session
+    if session_state is None or "session_id" not in session_state:
+        session_state = {"session_id": str(uuid.uuid4()), "page_count": 0}
+
+    session_manager = SessionManager(session_state["session_id"])
+
+    # Clean up old sessions periodically
+    if random.random() < 0.1:  # 10% chance to cleanup on each request
+        session_manager.cleanup_old_sessions()
+
+    # Check page limit
+    if session_manager.metadata["total_pages"] >= 128:
+        return session_state, None, None, "Maximum page limit (128) reached!", f"Page limit reached"
+
+    generated_images = []
+    used_seeds = []
+
+    # Generate the requested number of images
+    for i in range(int(num_images)):
+        progress(i / num_images, f"Generating image {i+1} of {num_images} for page {session_manager.metadata['total_pages'] + 1}")
+
+        current_seed = seed + i if not randomize_seed else random.randint(0, MAX_SEED)
+
+        # Generate single image
+        image, used_seed = infer_single(
+            prompt=prompt,
+            seed=current_seed,
+            randomize_seed=False,  # We handle randomization here
+            aspect_ratio=aspect_ratio,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            prompt_enhance=prompt_enhance,
+            style_preset=style_preset,
+            custom_style_text=custom_style_text,
+        )
+
+        generated_images.append(image)
+        used_seeds.append(used_seed)
+
+    # Add page to session
+    progress(0.8, "Adding page to document...")
+    page_num = session_manager.add_page(generated_images, layout, used_seeds)
+
+    # Create multi-page PDF
+    progress(0.9, "Creating PDF...")
+    pdf_path = create_multi_page_pdf(session_manager)
+
+    progress(1.0, "Done!")
+
+    # Update session state
+    session_state["page_count"] = page_num
+
+    # Prepare page info
+    seeds_str = ", ".join(str(s) for s in used_seeds)
+    page_info = f"Page {page_num} added\nSeeds: {seeds_str}\nTotal pages: {page_num}"
+
+    # Next button label
+    next_page_num = page_num + 1
+    button_label = f"Generate page {next_page_num}" if next_page_num <= 128 else "Page limit reached"
+
+    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(
+    prompt,
+    seed=42,
+    randomize_seed=False,
+    aspect_ratio="1:1",
+    guidance_scale=1.0,
+    num_inference_steps=8,
+    prompt_enhance=True,
+    style_preset="no_style",
+    custom_style_text="",
 ):
     """
     Generates an image based on a text prompt using the Qwen-Image-Lightning model.
@@ -181,42 +592,50 @@ def infer(
         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 
+        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 
+        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 
+        tuple[Image.Image, int]: A tuple containing the generated PIL Image and
             the integer seed used for the generation.
     """
-    # Use a blank negative prompt as per the lightning model's recommendation
-    negative_prompt = " "
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     # Convert aspect ratio to width and height
     width, height = get_image_size(aspect_ratio)
-    
+
     # Set up the generator for reproducibility
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    
-    print(f"Calling pipeline with prompt: '{prompt}'")
+
+    print(f"Original prompt: '{prompt}'")
+    print(f"Style preset: '{style_preset}'")
+
+    # Apply style preset first
+    styled_prompt, style_negative_prompt = apply_style_preset(prompt, style_preset, custom_style_text)
+
+    # Then apply prompt enhancement if enabled
     if prompt_enhance:
-        prompt = rewrite(prompt)
-        
-    print(f"Actual Prompt: '{prompt}'")
+        styled_prompt = rewrite(styled_prompt)
+
+    # Use style negative prompt if available, otherwise default
+    negative_prompt = style_negative_prompt if style_negative_prompt else " "
+
+    print(f"Final Prompt: '{styled_prompt}'")
     print(f"Negative Prompt: '{negative_prompt}'")
     print(f"Seed: {seed}, Size: {width}x{height}, Steps: {num_inference_steps}, True CFG Scale: {guidance_scale}")
 
     # Generate the image
     image = pipe(
-        prompt=prompt,
+        prompt=styled_prompt,
         negative_prompt=negative_prompt,
         width=width,
         height=height,
@@ -227,6 +646,9 @@ def infer(
 
     return image, seed
 
+# Keep the old infer function for backward compatibility
+infer = infer_single
+
 # --- Examples and UI Layout ---
 examples = [
         "A capybara wearing a suit holding a sign that reads Hello World",
@@ -254,6 +676,9 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
+    # Session state
+    session_state = gr.State(value={"session_id": str(uuid.uuid4()), "page_count": 0})
+
     with gr.Column(elem_id="col-container"):
         gr.HTML("""
         <div id="logo-title">
@@ -269,9 +694,18 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            run_button = gr.Button("Run", scale=0, variant="primary")
+            with gr.Column(scale=0):
+                run_button = gr.Button("Generate page 1", variant="primary")
+                reset_button = gr.Button("Start New Document", variant="secondary")
 
-        result = gr.Image(label="Result", show_label=False, type="pil")
+        with gr.Row():
+            with gr.Column(scale=1):
+                result_preview = gr.Image(label="Preview", show_label=True, type="pil")
+            with gr.Column(scale=1):
+                pdf_output = gr.File(label="Download PDF", show_label=True)
+                page_info = gr.Textbox(label="Page Info", show_label=True, interactive=False, lines=3)
+                gr.Markdown("""**Note:** Your images and PDF are saved for up to 24 hours.
+                You can continue adding pages (up to 128) by clicking the generate button.""")
 
         with gr.Accordion("Advanced Settings", open=False):
             seed = gr.Slider(
@@ -292,13 +726,48 @@ with gr.Blocks(css=css) as demo:
                 )
                 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)",
                     minimum=1.0,
                     maximum=5.0,
                     step=0.1,
-                    value=1.0, 
+                    value=1.0,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -309,11 +778,47 @@ with gr.Blocks(css=css) as demo:
                     value=8,
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
+            # 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]
+            )
+
+            # 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]
+            )
+
+        # 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],
+            ["A knight fighting a dragon", "medieval", "", 4],
+        ]
+
+        gr.Examples(
+            examples=styled_examples,
+            inputs=[prompt, style_preset, custom_style_text, num_images_slider],
+            outputs=None,  # Don't show outputs for examples
+            fn=None,
+            cache_examples=False
+        )
 
-    gr.on(
+    # Define the main generation event
+    generation_event = gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer,
+        fn=infer_page,
         inputs=[
             prompt,
             seed,
@@ -322,8 +827,25 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
             prompt_enhance,
+            style_preset,
+            custom_style_text,
+            num_images_slider,
+            layout_dropdown,
+            session_state,
         ],
-        outputs=[result, seed],
+        outputs=[session_state, pdf_output, result_preview, page_info, run_button],
+    )
+
+    # Reset button functionality
+    def reset_session():
+        new_state = {"session_id": str(uuid.uuid4()), "page_count": 0}
+        return new_state, None, None, "", "Generate page 1"
+
+    # Connect the reset button
+    reset_button.click(
+        fn=reset_session,
+        inputs=[],
+        outputs=[session_state, pdf_output, result_preview, page_info, run_button]
     )
 
 if __name__ == "__main__":

page_layouts.yaml

@@ -0,0 +1,127 @@
+# Page layouts configuration for multi-image PDF generation
+# Each layout defines how images are arranged on a page
+# Positions are defined as (x, y, width, height) in relative units (0-1)
+
+layouts:
+  1_image:
+    - id: "full_page"
+      label: "Full Page"
+      description: "Single image covering the full page"
+      positions:
+        - [0.05, 0.05, 0.9, 0.9]  # x, y, width, height (5% margins)
+
+  2_images:
+    - id: "horizontal_split"
+      label: "Layout A - Horizontal Split"
+      description: "Two images side by side"
+      positions:
+        - [0.05, 0.05, 0.425, 0.9]  # Left image
+        - [0.525, 0.05, 0.425, 0.9]  # Right image
+
+    - id: "vertical_split"
+      label: "Layout B - Vertical Split"
+      description: "Two images stacked vertically"
+      positions:
+        - [0.05, 0.05, 0.9, 0.425]  # Top image
+        - [0.05, 0.525, 0.9, 0.425]  # Bottom image
+
+    - id: "dominant_left"
+      label: "Layout C - Large Left"
+      description: "Large image on left, small on right"
+      positions:
+        - [0.05, 0.05, 0.6, 0.9]  # Large left image
+        - [0.7, 0.25, 0.25, 0.5]  # Small right image
+
+    - id: "dominant_top"
+      label: "Layout D - Large Top"
+      description: "Large image on top, small on bottom"
+      positions:
+        - [0.05, 0.05, 0.9, 0.6]  # Large top image
+        - [0.25, 0.7, 0.5, 0.25]  # Small bottom image
+
+  3_images:
+    - id: "grid_horizontal"
+      label: "Layout A - Horizontal Strip"
+      description: "Three images in a row"
+      positions:
+        - [0.05, 0.25, 0.283, 0.5]  # Left
+        - [0.358, 0.25, 0.283, 0.5]  # Middle
+        - [0.666, 0.25, 0.283, 0.5]  # Right
+
+    - id: "grid_vertical"
+      label: "Layout B - Vertical Strip"
+      description: "Three images in a column"
+      positions:
+        - [0.25, 0.05, 0.5, 0.283]  # Top
+        - [0.25, 0.358, 0.5, 0.283]  # Middle
+        - [0.25, 0.666, 0.5, 0.283]  # Bottom
+
+    - id: "hero_top"
+      label: "Layout C - Hero Top"
+      description: "Large image on top, two small below"
+      positions:
+        - [0.05, 0.05, 0.9, 0.5]  # Large top
+        - [0.05, 0.6, 0.425, 0.35]  # Bottom left
+        - [0.525, 0.6, 0.425, 0.35]  # Bottom right
+
+    - id: "hero_left"
+      label: "Layout D - Hero Left"
+      description: "Large image on left, two small on right"
+      positions:
+        - [0.05, 0.05, 0.5, 0.9]  # Large left
+        - [0.6, 0.05, 0.35, 0.425]  # Top right
+        - [0.6, 0.525, 0.35, 0.425]  # Bottom right
+
+    - id: "diagonal"
+      label: "Layout E - Diagonal"
+      description: "Diagonal arrangement"
+      positions:
+        - [0.05, 0.05, 0.4, 0.4]  # Top left
+        - [0.3, 0.3, 0.4, 0.4]  # Center
+        - [0.55, 0.55, 0.4, 0.4]  # Bottom right
+
+  4_images:
+    - id: "grid_2x2"
+      label: "Layout A - 2x2 Grid"
+      description: "Four equal images in a grid"
+      positions:
+        - [0.05, 0.05, 0.425, 0.425]  # Top left
+        - [0.525, 0.05, 0.425, 0.425]  # Top right
+        - [0.05, 0.525, 0.425, 0.425]  # Bottom left
+        - [0.525, 0.525, 0.425, 0.425]  # Bottom right
+
+    - id: "strip_horizontal"
+      label: "Layout B - Horizontal Strip"
+      description: "Four images in a row"
+      positions:
+        - [0.05, 0.3, 0.2125, 0.4]  # First
+        - [0.2875, 0.3, 0.2125, 0.4]  # Second
+        - [0.525, 0.3, 0.2125, 0.4]  # Third
+        - [0.7625, 0.3, 0.2125, 0.4]  # Fourth
+
+    - id: "strip_vertical"
+      label: "Layout C - Vertical Strip"
+      description: "Four images in a column"
+      positions:
+        - [0.3, 0.05, 0.4, 0.2125]  # First
+        - [0.3, 0.2875, 0.4, 0.2125]  # Second
+        - [0.3, 0.525, 0.4, 0.2125]  # Third
+        - [0.3, 0.7625, 0.4, 0.2125]  # Fourth
+
+    - id: "hero_with_strip"
+      label: "Layout D - Hero with Strip"
+      description: "One large image with three small ones"
+      positions:
+        - [0.05, 0.05, 0.6, 0.6]  # Large main
+        - [0.7, 0.05, 0.25, 0.283]  # Small top
+        - [0.7, 0.358, 0.25, 0.283]  # Small middle
+        - [0.7, 0.666, 0.25, 0.283]  # Small bottom
+
+    - id: "l_shape"
+      label: "Layout E - L Shape"
+      description: "L-shaped arrangement"
+      positions:
+        - [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)

requirements.txt

@@ -4,4 +4,7 @@ transformers
 accelerate
 safetensors
 sentencepiece
-dashscope
+dashscope
+pyyaml
+reportlab
+pypdf2

style_presets.yaml

@@ -0,0 +1,221 @@
+# Style presets for image generation
+# Each preset defines how to modify the user's prompt with style-specific prefixes and suffixes
+
+presets:
+  no_style:
+    id: "no_style"
+    label: "No style (custom)"
+    enabled: true
+    prompt_prefix: ""
+    prompt_suffix: ""
+    negative_prompt: ""
+
+  random:
+    id: "random"
+    label: "Random style"
+    enabled: true
+    # This will be handled specially in code
+    prompt_prefix: ""
+    prompt_suffix: ""
+    negative_prompt: ""
+
+  japanese_manga:
+    id: "japanese_manga"
+    label: "Japanese Manga"
+    enabled: true
+    prompt_prefix: "grayscale, detailed drawing, japanese manga"
+    prompt_suffix: ""
+    negative_prompt: "franco-belgian comic, color album, color, american comic, photo, painting, 3D render"
+
+  nihonga:
+    id: "nihonga"
+    label: "Nihonga"
+    enabled: true
+    prompt_prefix: "japanese nihonga painting about"
+    prompt_suffix: "Nihonga, ancient japanese painting, intricate, detailed, detailed painting"
+    negative_prompt: "franco-belgian comic, color album, color, manga, comic, american comic, photo, painting, 3D render"
+
+  franco_belgian:
+    id: "franco_belgian"
+    label: "Franco-Belgian"
+    enabled: true
+    prompt_prefix: "bande dessinée, franco-belgian comic"
+    prompt_suffix: "comic album, detailed drawing"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, photo, painting, 3D render"
+
+  american_comic_90:
+    id: "american_comic_90"
+    label: "American Comic (modern)"
+    enabled: true
+    prompt_prefix: "digital color comicbook style, modern american comic"
+    prompt_suffix: "detailed drawing"
+    negative_prompt: "manga, anime, american comic, action, grayscale, monochrome, photo, painting, 3D render"
+
+  # american_comic_40 is commented out in the original
+  # american_comic_40:
+  #   id: "american_comic_40"
+  #   label: "American Comic (1940)"
+  #   enabled: false
+  #   prompt_prefix: "american comic"
+  #   prompt_suffix: "single panel, american comic, comicbook style, 1940, 40s, color comicbook, color drawing"
+  #   negative_prompt: "manga, anime, american comic, action, grayscale, monochrome, photo, painting, 3D render"
+
+  american_comic_50:
+    id: "american_comic_50"
+    label: "American Comic (1950)"
+    enabled: true
+    prompt_prefix: "1950, 50s, vintage american color comic"
+    prompt_suffix: "detailed drawing"
+    negative_prompt: "manga, anime, american comic, action, grayscale, monochrome, photo, painting, 3D render"
+
+  # american_comic_60 is commented out in the original
+  # american_comic_60:
+  #   id: "american_comic_60"
+  #   label: "American Comic (1960)"
+  #   enabled: false
+  #   prompt_prefix: "american comic"
+  #   prompt_suffix: "single panel, american comic, comicbook style, 1960, 60s, color comicbook, color drawing"
+  #   negative_prompt: "manga, anime, american comic, action, grayscale, monochrome, photo, painting, 3D render"
+
+  flying_saucer:
+    id: "flying_saucer"
+    label: "Flying Saucer (Pulp Sci-Fi)"
+    enabled: true
+    prompt_prefix: "vintage science fiction, color pulp comic panel, 1940"
+    prompt_suffix: "detailed drawing"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, photo, painting, 3D render"
+
+  humanoid:
+    id: "humanoid"
+    label: "Humanoid (Moebius)"
+    enabled: true
+    prompt_prefix: "color comic panel, style of Moebius"
+    prompt_suffix: "detailed drawing, french comic panel, franco-belgian style, bande dessinée, single panel"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, photo, painting, 3D render"
+
+  haddock:
+    id: "haddock"
+    label: "Haddock (Tintin)"
+    enabled: true
+    prompt_prefix: "color comic panel, style of Hergé, tintin style"
+    prompt_suffix: "by Hergé, french comic panel, franco-belgian style"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, photo, painting, 3D render"
+
+  # lurid is commented out in the original
+  # lurid:
+  #   id: "lurid"
+  #   label: "Lurid (Underground)"
+  #   enabled: false
+  #   prompt_prefix: "satirical color comic, underground comix, 1970"
+  #   prompt_suffix: ""
+  #   negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  armorican:
+    id: "armorican"
+    label: "Armorican (Asterix)"
+    enabled: true
+    prompt_prefix: "color comic panel, romans, gauls, french comic panel, franco-belgian style, about"
+    prompt_suffix: "bande dessinée, single panel"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, photo, painting, 3D render"
+
+  render:
+    id: "render"
+    label: "3D Render (Pixar)"
+    enabled: true
+    prompt_prefix: "3D render animation, Pixar, cute, funny, Unreal engine"
+    prompt_suffix: "crisp, sharp"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  klimt:
+    id: "klimt"
+    label: "Klimt"
+    enabled: true
+    prompt_prefix: "golden, patchwork, style of Gustav Klimt, Gustav Klimt painting"
+    prompt_suffix: "detailed painting, intricate details"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  medieval:
+    id: "medieval"
+    label: "Medieval Illuminated"
+    enabled: true
+    prompt_prefix: "medieval illuminated manuscript, illuminated manuscript of, medieval"
+    prompt_suffix: "intricate details"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  # glass is commented out in the original
+  # glass:
+  #   id: "glass"
+  #   label: "Stained Glass"
+  #   enabled: false
+  #   prompt_prefix: "stained glass, vitrail, stained glass"
+  #   prompt_suffix: "medieval"
+  #   negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  # voynich is commented out in the original
+  # voynich:
+  #   id: "voynich"
+  #   label: "Voynich"
+  #   enabled: false
+  #   prompt_prefix: "voynich, voynich page"
+  #   prompt_suffix: "medieval"
+  #   negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  egyptian:
+    id: "egyptian"
+    label: "Egyptian"
+    enabled: true
+    prompt_prefix: "ancient egyptian wall painting, ancient egypt"
+    prompt_suffix: ""
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  # psx is commented out in the original
+  # psx:
+  #   id: "psx"
+  #   label: "PSX"
+  #   enabled: false
+  #   prompt_prefix: "videogame screenshot, 3dfx, 3D dos game, software rendering"
+  #   prompt_suffix: ""
+  #   negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  # pixel is commented out in the original
+  # pixel:
+  #   id: "pixel"
+  #   label: "Pixel Art"
+  #   enabled: false
+  #   prompt_prefix: "pixelart, isometric, pixelated, low res"
+  #   prompt_suffix: ""
+  #   negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  photonovel:
+    id: "photonovel"
+    label: "Vintage Photonovel"
+    enabled: true
+    prompt_prefix: "vintage photo, 1950, 1960, french new wave, faded colors, color movie screencap"
+    prompt_suffix: ""
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+  stockphoto:
+    id: "stockphoto"
+    label: "Stock Photo"
+    enabled: true
+    prompt_prefix: "cinematic, hyperrealistic, footage, sharp 8k, analog, instagram, photoshoot"
+    prompt_suffix: "crisp details"
+    negative_prompt: "manga, anime, american comic, grayscale, monochrome, painting"
+
+# video_3d_style was in the original but commented out
+# video_3d_style:
+#   id: "video_3d_style"
+#   label: "[video] 3D style"
+#   enabled: false
+#   prompt_prefix: ""
+#   prompt_suffix: ""
+#   negative_prompt: ""
+
+# neutral preset from original
+  neutral:
+    id: "neutral"
+    label: "Neutral"
+    enabled: true
+    prompt_prefix: ""
+    prompt_suffix: ""
+    negative_prompt: ""