orchestrator.py

import base64
import io
import json
import os
import sys
import textwrap
from typing import Dict

from PIL import Image, ImageDraw, ImageFont

import modal
from pydantic import BaseModel, Field, ValidationError

sys.path.append("/root")
from constants import (
    CHARACTER_ANCHOR,
    GLOBAL_NEGATIVE_PROMPT,
    GLOBAL_STYLE_MODIFIERS,
)

# Create primary orchestrator app wrapper.
image = (
    modal.Image.debian_slim()
    .apt_install("fontconfig", "fonts-liberation")
    .pip_install(
        "accelerate",
        "diffusers",
        "fastapi[standard]",
        "huggingface_hub>=0.23.0",
        "pillow",
        "pydantic>=2.0.0",
        "torch",
        "transformers",
    )
    .add_local_file("constants.py", "constants.py")
    .add_local_file("comic.ttf", "comic.ttf")
)

app = modal.App("perri-comic-pipeline", image=image)


# =====================================================================
# SHARED SCHEMAS AND PROMPTS
# =====================================================================
MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"


class ComicPanel(BaseModel):
    panel_number: int = Field(
        ...,
        description="The sequential integer of the panel, starting at 1.",
    )
    visual_description: str = Field(
        ...,
        description="Granular visual actions and environmental changes only. Do not mention text or speech bubbles here.",
    )
    dialogue: str = Field(
        ...,
        description="The exact speech bubble text or narrative caption for this panel.",
    )


class ComicPanelScript(BaseModel):
    panel: ComicPanel = Field(
        ...,
        description="The complete single comic panel.",
    )


SYSTEM_PROMPT = (
    "You are a strict, non-conversational comic book continuity scriptwriter.\n"
    "Your job is to break the user's storyline down into a sequential panel-by-panel script.\n"
    "CRITICAL: You must output ONLY the requested JSON structure conforming exactly to the schema provided. "
    "Do not include introductory text, conversational filler, markdown formatting blocks (like ```json), or sign-offs. "
    "Focus heavily on blocking physical action, physical movement, and dramatic environmental modifications panel by panel.\n"
    "Never repeat the same word or phrase more than 3 times consecutively.\n"
    "Summarize repetitive actions or sounds, e.g. '(meows repeatedly)' instead of repeating 'Meow'.\n"
    "Keep dialogue concise."
)


# =====================================================================
# HELPER FUNCTIONS
# =====================================================================

def _strip_code_fences(raw_text: str) -> str:
    """Removes standard or JSON markdown code block formatting backticks from strings."""
    cleaned_text = raw_text.strip()
    if cleaned_text.startswith("```json"):
        return cleaned_text.split("```json", 1)[1].split("```", 1)[0].strip()
    if cleaned_text.startswith("```"):
        return cleaned_text.split("```", 1)[1].split("```", 1)[0].strip()
    return cleaned_text


def _load_dialogue_font(size: int = 28) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    """Load the bundled comic font when available, otherwise fall back cleanly."""
    font_candidates = (
        "comic.ttf",
        "/root/comic.ttf",
        "/assets/comic.ttf",
    )

    for font_path in font_candidates:
        try:
            return ImageFont.truetype(font_path, size=size)
        except OSError:
            continue

    return ImageFont.load_default()


def overlay_dialogue_bubble(image_bytes: bytes, dialogue_text: str) -> bytes:
    """
    Renders a clean dialogue box containing standard text over an image.
    Uses hardcoded bitmap character dimensions to prevent layout collapses 
    in stateless containerized execution environments like Hugging Face Spaces.
    """
    # Open image data and guarantee an alpha channel for standard drawing properties
    base_img = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
    width, height = base_img.size
    
    # Clean text inputs and break lines appropriately using textwrap
    clean_text = _strip_code_fences(dialogue_text)
    font = _load_dialogue_font(size=35)

    # Dynamically wrap based on image width so the text remains readable on both
    # small and large panels.
    wrap_width = max(18, min(34, width // 18))
    wrapped_lines = textwrap.wrap(clean_text, width=wrap_width)
    if not wrapped_lines:
        wrapped_lines = [""]
    sample_text = "\n".join(wrapped_lines)

    # Measure text using the real font so the bubble scales correctly.
    layout_probe = Image.new("RGBA", (width, height))
    layout_draw = ImageDraw.Draw(layout_probe)
    bbox = layout_draw.multiline_textbbox((0, 0), sample_text, font=font, spacing=8, stroke_width=2)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]

    margin = 24
    padding_x = 24
    padding_y = 18

    box_left = margin
    box_right = min(width - margin, box_left + text_width + (padding_x * 2))
    box_bottom = height - margin
    box_top = max(margin, box_bottom - text_height - (padding_y * 2))

    if (box_right - box_left) < 180:
        box_right = min(width - margin, box_left + 180)
    if (box_bottom - box_top) < 90:
        box_top = max(margin, box_bottom - 90)
    
    # Setup rendering canvas wrapper
    final_img = base_img.copy()
    draw = ImageDraw.Draw(final_img)
    
    # Render a high-contrast speech bubble so the dialogue is visible on noisy art.
    draw.rounded_rectangle(
        [box_left, box_top, box_right, box_bottom],
        radius=16,
        fill=(255, 248, 228, 242),
        outline=(30, 24, 18, 255),
        width=4,
    )
    
    # Overlay lines onto image canvas matching coordinates calculated above
    text_x = box_left + padding_x
    text_y = box_top + padding_y
    draw.multiline_text(
        (text_x, text_y),
        sample_text,
        fill=(20, 16, 12, 255),
        font=font,
        spacing=8,
        stroke_width=2,
        stroke_fill=(255, 255, 255, 255),
    )
    
    # Flatten composite canvas down to pure RGB matrix and dump to JPEG stream
    output_stream = io.BytesIO()
    final_img.convert("RGB").save(output_stream, format="JPEG", quality=90)
    return output_stream.getvalue()


def _error_payload(message: str) -> Dict:
    """Return a frontend-friendly error body without raising a 500."""
    return {
        "status": "error",
        "message": message,
        "panel": {
            "id": 1,
            "image_b64": "",
        },
        "script": None,
    }
# =====================================================================
# WORKER FUNCTIONS
# =====================================================================
@app.function(image=image, secrets=[modal.Secret.from_name("huggingface-secret")])
def generate_script(user_idea: str, theme: str, tone: str) -> str:
    from huggingface_hub import InferenceClient

    client = InferenceClient(
        model=MODEL_NAME,
        api_key=os.environ["HF_TOKEN"],
    )

    user_instructions = (
        f"Story Outline: {user_idea}\n"
        f"Core Theme Formula: {theme}\n"
        f"Narrative Tone: {tone}\n"
        "Total Panels Requested: 1"
    )

    response_format = {
        "type": "json_schema",
        "json_schema": {
            "name": "ComicPanelScript",
            "schema": ComicPanelScript.model_json_schema(),
            "strict": True,
        },
    }

    completion = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_instructions},
        ],
        max_tokens=1500,
        temperature=0.7,
        frequency_penalty=1.0,
        response_format=response_format,
    )

    raw_json_string = completion.choices[0].message.content
    if raw_json_string is None:
        raise RuntimeError("Hugging Face returned an empty response body.")
    return raw_json_string


@app.function(image=image, gpu="A10G", timeout=60)
def generate_panel_image(payload: dict) -> bytes:
    from diffusers import AutoPipelineForText2Image
    import torch

    panel_action = payload.get("action", "")
    color_mode = payload.get("color_mode", "Black and White")

    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=torch.float16,
        variant="fp16",
    ).to("cuda")

    style_modifiers = GLOBAL_STYLE_MODIFIERS
    if color_mode == "Black and White":
        style_modifiers += ", retro monochrome comic book print, black and white ink sketch, no color"
    else:
        style_modifiers += ", rich vibrant color palette"

    final_prompt = f"{CHARACTER_ANCHOR} Action: {panel_action}, {style_modifiers}"

    image_object = pipe(
        prompt=final_prompt,
        negative_prompt=GLOBAL_NEGATIVE_PROMPT,
        num_inference_steps=2,
        guidance_scale=0.0,
    ).images[0]

    byte_stream = io.BytesIO()
    image_object.save(byte_stream, format="JPEG", quality=85)
    return byte_stream.getvalue()


# =====================================================================
# UNIFIED API GATEWAY
# =====================================================================
@app.function(image=image)
@modal.fastapi_endpoint(method="POST")
def generate_comic_api(request_data: Dict) -> Dict:
    try:
        if not isinstance(request_data, dict):
            raise ValueError("Request body must be a JSON object.")

        user_prompt = str(request_data.get("prompt", "")).strip()
        color_mode = str(request_data.get("color_mode", "Black and White")).strip()
        theme = str(request_data.get("theme", "Potion")).strip()
        tone = str(request_data.get("tone", "")).strip()

        print(f"Processing request for story seed: '{user_prompt}' (1 panel)")

        raw_script_json = generate_script.remote(
            user_idea=user_prompt,
            theme=theme,
            tone=tone,
        )

        cleaned_json_string = _strip_code_fences(raw_script_json)

        try:
            script_data = ComicPanelScript.model_validate_json(cleaned_json_string)
        except ValidationError as ve:
            print(f"LLM returned invalid raw string: {raw_script_json}")
            return _error_payload(f"Failed to validate LLM script response: {ve}")

        panel_data = script_data.panel.model_dump()

        print("Launching GPU worker for single-panel generation...")
        raw_image_bytes = generate_panel_image.remote(
            {
                "action": panel_data["visual_description"],
                "color_mode": color_mode,
            }
        )

        print("Compositing script dialogue bubble into panel...")
        finished_panel_bytes = overlay_dialogue_bubble(
            raw_image_bytes,
            panel_data.get("dialogue", ""),
        )
        base64_string = base64.b64encode(finished_panel_bytes).decode("utf-8")

        output_payload = {
            "status": "success",
            "script": panel_data,
            "panel": {
                "id": 1,
                "image_b64": base64_string,
            },
        }

        print("Comic strip processing complete. Shipping payload back to frontend!")
        return output_payload
    except Exception as exc:
        print(f"Orchestrator failed: {exc}")
        return _error_payload(str(exc))


@app.local_entrypoint()
def main() -> None:
    sample_prompt = "A young knight wants an elixir to look 20 years younger, but it turns him into a toddler."
    print("Launching Modal container to orchestrate structured text pipeline...")

    json_output = generate_script.remote(
        user_idea=sample_prompt,
        theme="Mortal Request / Potion Twist",
        tone="Humor",
    )

    print("\nClean JSON Structure Returned Successfully:")
    print(json_output)