text_to_story.py

import json
import os
import io
import requests
import textwrap
from dotenv import load_dotenv
from PIL import Image, ImageDraw, ImageFont
import replicate
import time

# Load the .env file
load_dotenv()

# Set up Azure OpenAI endpoint and key
AZURE_OPENAI_KEY = os.getenv("azure_open_ai_key")
AZURE_ENDPOINT_URL = os.getenv("azure_endpoint_url")

# Set up Replicate API token
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

def generate_panels(scenario):
    if not AZURE_OPENAI_KEY or not AZURE_ENDPOINT_URL:
        raise ValueError("Azure OpenAI key or endpoint URL is not set. Please check your .env file.")

    prompt = f"""
    You are a comic artist creating a 6-panel comic strip. Given the following scenario, create a sequential story across 6 panels. For each panel, provide:
    1. A brief description of the scene (characters, background, and camera angle)
    2. The dialogue or text for that panel

    Use a variety of camera angles to make the comic visually interesting. Choose from these angles:
    - Close-up: Focuses on a character's face or a specific object
    - Medium shot: Shows characters from the waist up
    - Wide shot: Shows the entire scene, including all characters and the environment
    - Over-the-shoulder: Views the scene from behind a character, looking at what they're seeing
    - Bird's-eye view: Looks down on the scene from above
    - Low angle: Looks up at the characters or scene from below

    Ensure that the story progresses logically from panel to panel, showing the characters' journey from the initial idea to the final presentation.

    Scenario:
    {scenario}

    Format your response as follows for each panel:
    Panel 1:
    Description: [Camera angle] - [Scene description]
    Text: [Character name]: "[Dialogue]"

    Panel 2:
    ...

    Remember to maintain continuity and show progression in the story across all 6 panels while varying the camera angles.
    """

    try:
        headers = {
            "Content-Type": "application/json",
            "api-key": AZURE_OPENAI_KEY
        }

        payload = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant that creates sequential comic strips with varied camera angles."},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7,
            "max_tokens": 1000
        }

        response = requests.post(AZURE_ENDPOINT_URL, headers=headers, json=payload)
        
        if response.status_code != 200:
            print(f"Error response from API: {response.status_code}")
            print(f"Response content: {response.text}")
            return []

        result = response.json()
        generated_text = result['choices'][0]['message']['content']
        print(generated_text)

        return extract_panel_info(generated_text)

    except Exception as e:
        print(f"An error occurred in generate_panels: {str(e)}")
        return []

def extract_panel_info(text):
    panel_info_list = []
    panels = text.split('Panel')[1:]  # Split by 'Panel' and remove the first empty element

    for i, panel in enumerate(panels, start=1):
        panel_info = {}
        lines = panel.strip().split('\n')
        
        panel_info['number'] = str(i)
        
        # Extract description with camera angle
        description_line = next((line for line in lines if line.strip().startswith('Description:')), '')
        if description_line:
            parts = description_line.split('-', 1)
            if len(parts) == 2:
                panel_info['camera_angle'] = parts[0].split(':', 1)[1].strip()
                panel_info['description'] = parts[1].strip()
            else:
                panel_info['camera_angle'] = 'Not specified'
                panel_info['description'] = description_line.split(':', 1)[1].strip()
        else:
            panel_info['camera_angle'] = 'Not specified'
            panel_info['description'] = ''
        
        # Extract text
        text_line = next((line for line in lines if line.strip().startswith('Text:')), '')
        panel_info['text'] = text_line.split(':', 1)[1].strip() if text_line else ''

        panel_info_list.append(panel_info)

    return panel_info_list

def generate_image(prompt, max_retries=3, retry_delay=5):
    for attempt in range(max_retries):
        try:
            output = replicate.run(
                "dharmagnavyas/sdxl-shortfilm:e57f0bc20475e364dd74631d602a069b1cd5eff7d93d735aa4f6ef7d53ea1c77",
                input={
                    "width": 1024,
                    "height": 1024,
                    "prompt": prompt,
                    "refine": "no_refiner",
                    "scheduler": "K_EULER",
                    "lora_scale": 0.6,
                    "num_outputs": 1,
                    "guidance_scale": 7.5,
                    "apply_watermark": True,
                    "high_noise_frac": 0.8,
                    "negative_prompt": "",
                    "prompt_strength": 0.8,
                    "num_inference_steps": 50
                },
                timeout=90  # Set a 90-second timeout
            )
            
            if output and isinstance(output, list) and len(output) > 0:
                image_url = output[0]
                response = requests.get(image_url, timeout=30)  # Set a 30-second timeout for downloading the image
                img = Image.open(io.BytesIO(response.content))
                return img
            else:
                print(f"Attempt {attempt + 1}: Unexpected output format from Replicate API")
        except requests.exceptions.Timeout:
            print(f"Attempt {attempt + 1}: The request to Replicate API timed out")
        except Exception as e:
            print(f"Attempt {attempt + 1}: An error occurred in generate_image: {str(e)}")
        
        if attempt < max_retries - 1:
            print(f"Retrying in {retry_delay} seconds...")
            time.sleep(retry_delay)
    
    print(f"Failed to generate image after {max_retries} attempts")
    return None

def add_text_to_panel(text, panel_image):
    text_image_height = 150
    text_image = generate_text_image(text, panel_image.width, text_image_height)

    result_image = Image.new('RGB', (panel_image.width, panel_image.height + text_image_height))
    result_image.paste(panel_image, (0, 0))
    result_image.paste(text_image, (0, panel_image.height))

    return result_image

def generate_text_image(text, width, height):
    image = Image.new('RGB', (width, height), color='white')
    draw = ImageDraw.Draw(image)

    font = ImageFont.truetype(font="manga-temple.ttf", size=28)

    wrapped_text = textwrap.wrap(text, width=35)
    
    y_text = 10
    for line in wrapped_text:
        left, top, right, bottom = draw.textbbox((0, 0), line, font=font)
        text_width = right - left
        x = (width - text_width) // 2
        draw.text((x, y_text), line, font=font, fill="black")
        y_text += bottom - top + 8

    return image

def create_strip(images):
    columns, rows = 2, 3
    padding = 20

    panel_width = images[0].width
    panel_height = images[0].height
    output_width = columns * panel_width + (columns + 1) * padding
    output_height = rows * panel_height + (rows + 1) * padding

    result_image = Image.new("RGB", (output_width, output_height), "white")

    for i, img in enumerate(images):
        x = padding + (i % columns) * (panel_width + padding)
        y = padding + (i // columns) * (panel_height + padding)
        result_image.paste(img, (x, y))

    return result_image.resize((1024, int(1024 * output_height / output_width)))

def generate_storyboard(scenario):
    print(f"Generating panels for this scenario: \n {scenario}")

    try:
        panels = generate_panels(scenario)
        if not panels:
            print("No panels were generated. Please check your Azure OpenAI service configuration.")
            return None
        else:
            print(f"Generated panels: {json.dumps(panels, indent=2)}")

            os.makedirs('output2', exist_ok=True)
            
            with open('output/panels.json', 'w') as outfile:
                json.dump(panels, outfile, indent=2)

            panel_images = []

            for panel in panels:
                panel_prompt = f"{panel['camera_angle']} view - {panel['description']}"
                print(f"Generate panel {panel['number']} with prompt: {panel_prompt}")
                panel_image = generate_image(panel_prompt)
                if panel_image:
                    panel_image_with_text = add_text_to_panel(panel["text"], panel_image)
                    panel_image_with_text.save(f"output/panel-{panel['number']}.png")
                    panel_images.append(panel_image_with_text)
                else:
                    print(f"Failed to generate image for panel {panel['number']}")

            if panel_images:
                try:
                    final_strip = create_strip(panel_images)
                    final_strip.save("output/strip.png")
                    print("Comic strip created successfully!")
                    return final_strip
                except Exception as e:
                    print(f"Error creating final comic strip: {str(e)}")
                    return None
            else:
                print("No panel images were generated. Unable to create the final comic strip.")
                return None

    except ValueError as ve:
        print(f"Configuration error: {str(ve)}")
        return None
    except Exception as e:
        print(f"An unexpected error occurred: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

if __name__ == "__main__":
    SCENARIO = """
Panel 1: Panoramic view of a desolate, futuristic cityscape. Crumbling skyscrapers covered in bioluminescent vines, hovering vehicles weaving through the ruins.
Text: Narrator: "TOK year is 2898. Earth as we knew it is long gone."
Panel 2: Close-up of our protagonist, a young woman with cybernetic enhancements, standing on a rooftop overlooking the city.
Text: Protagonist: "TOK ancestors called this place Mumbai. Now it's just Sector 7."
Panel 3: The protagonist enters a hidden lab filled with advanced technology and ancient artifacts.
Text: Scientist: "Asha, TOK found it! The key to saving our world!"
Panel 4: A holographic display shows a glowing artifact and data streams.
Text: Asha: "TOK Cosmic Seed? But that's just a myth from the old world."
Panel 5: Wide shot of a massive, menacing fortress in the distance, surrounded by energy fields.
Text: Scientist: "TOK have to get it before the Technocrats do. They'll use it to enslave what's left of humanity."
Panel 6: Asha and a small team gearing up with futuristic weapons and stealth suits.
Text: Asha: "TOK won't let that happen. It's time to take back our future.    """
    generate_storyboard(SCENARIO)