| from google.generativeai import GenerativeModel |
| import json |
| import re |
| import os |
| import datetime |
| import openai |
| import config |
| import time |
| from datetime import datetime |
|
|
|
|
| def log_execution(func): |
| def wrapper(*args, **kwargs): |
| start_time = time.time() |
| start_str = datetime.fromtimestamp(start_time).strftime('%Y-%m-%d %H:%M:%S') |
| |
| result = func(*args, **kwargs) |
| |
| end_time = time.time() |
| end_str = datetime.fromtimestamp(end_time).strftime('%Y-%m-%d %H:%M:%S') |
| duration = end_time - start_time |
|
|
| |
| return result |
| return wrapper |
|
|
| class StoryGenerator: |
| """ |
| Direct story generator that creates comic panel style stories from user input. |
| """ |
|
|
| def __init__(self): |
| self.model = GenerativeModel('gemini-2.5-flash') |
| @log_execution |
| def log_prompt(self, prompt, log_file="story_prompt_logs.jsonl"): |
| """Log the prompt to a file for debugging and improvement purposes.""" |
| log_entry = { |
| "timestamp": datetime.datetime.now().isoformat(), |
| "prompt": prompt |
| } |
| with open(log_file, "a", encoding="utf-8") as f: |
| f.write(json.dumps(log_entry, ensure_ascii=False) + "\n") |
| @log_execution |
| def enhance_user_story(self, user_description, max_retries=3, current_retry=0): |
| """ |
| Enhance the user's story with more vibrancy, detail, and narrative richness using |
| optimized AI prompting techniques for visual storytelling with smart detail preservation. |
| |
| Args: |
| user_description: The user's original story idea or prompt |
| max_retries: Maximum number of retry attempts (default: 3) |
| current_retry: Current retry attempt number (default: 0) |
| |
| Returns: |
| enhanced_story: A more vibrant and detailed version of the story with preserved key elements |
| """ |
| print(f"[StoryGenerator] Enhancing user story (attempt {current_retry + 1}/{max_retries}): {user_description[:100]}...") |
| |
| if current_retry >= max_retries: |
| print(f"[StoryGenerator] Max retries reached, returning original description") |
| return user_description |
| |
| try: |
| enhancement_prompt = self._create_detail_focused_enhancement_prompt(user_description) |
| |
| self.log_prompt(enhancement_prompt) |
| |
| try: |
| response = self.model.generate_content(enhancement_prompt) |
| enhanced_story = response.text.strip() |
| |
| if self._validate_enhancement_quality(enhanced_story, user_description): |
| print(f"[StoryGenerator] Story successfully enhanced with detail preservation") |
| return enhanced_story |
| else: |
| print(f"[StoryGenerator] Enhancement quality insufficient, using original with minimal enhancement") |
| return self._create_minimal_enhancement(user_description) |
| |
| except Exception as gemini_error: |
| print(f"[StoryGenerator] Gemini API error: {gemini_error}") |
| if current_retry < max_retries - 1: |
| print(f"[StoryGenerator] Retrying with simplified approach...") |
| return self._simplified_enhancement(user_description) |
| else: |
| raise gemini_error |
| |
| except Exception as e: |
| print(f"[StoryGenerator] Enhancement error: {e}") |
| if current_retry < max_retries - 1: |
| import time |
| time.sleep(1 * (current_retry + 1)) |
| return self.enhance_user_story(user_description, max_retries, current_retry + 1) |
| else: |
| print(f"[StoryGenerator] All enhancement attempts failed, returning original") |
| return user_description |
| @log_execution |
|
|
| def _create_detail_focused_enhancement_prompt(self, user_description): |
| """Create a concise enhancement prompt that adds coherence and enough detail for the required number of scenes.""" |
| return f""" |
| You are an expert visual storytelling assistant. Enhance the user's story concept to create a rich visual narrative. |
| |
| ORIGINAL STORY: "{user_description}" |
| |
| ENHANCEMENT GOALS: |
| • Define key character appearances (visual features, clothing). |
| • Establish a clear setting and atmosphere. |
| • Outline a logical scene progression that can be broken down into multiple action-focused panels. |
| • Ensure visual consistency for characters and locations. |
| • Descriptions should be concise yet vivid, focusing on elements crucial for an action-oriented digital comic. |
| |
| OUTPUT: Enhanced story description (2-3 paragraphs maximum) that provides a strong foundation for a multi-panel, action-focused visual story. Ensure the tone is suitable for a modern digital comic. |
| """ |
|
|
| def _validate_enhancement_quality(self, enhanced_story, original_story): |
| """Validate that the enhancement adds coherence and appropriate detail.""" |
| if not enhanced_story or len(enhanced_story) < 50: |
| return False |
| |
| enhanced_words = len(enhanced_story.split()) |
| original_words = len(original_story.split()) |
| |
| if enhanced_words < original_words or enhanced_words > original_words * 5: |
| return False |
| |
| story_elements = ['character', 'scene', 'story', 'visual', 'setting', 'action'] |
| has_story_elements = sum(1 for element in story_elements if element.lower() in enhanced_story.lower()) |
| |
| if has_story_elements < 2: |
| return False |
| |
| similarity_threshold = 0.8 |
| original_lower = original_story.lower() |
| enhanced_lower = enhanced_story.lower() |
| |
| common_words = set(original_lower.split()) & set(enhanced_lower.split()) |
| original_unique = len(set(original_lower.split())) |
| |
| if original_unique > 0: |
| similarity = len(common_words) / original_unique |
| if similarity > similarity_threshold and enhanced_words < original_words * 1.5: |
| return False |
| |
| return True |
| @log_execution |
|
|
| def _create_minimal_enhancement(self, user_description): |
| """Create minimal enhancement that preserves original while adding basic coherence for the required number of scenes.""" |
| |
| enhanced = f""" |
| Enhanced Story: {user_description} |
| |
| Visual Coherence Elements: |
| - Main character with consistent appearance throughout all scenes |
| - Clear setting that remains visually consistent |
| - Logical progression suitable for the required number of sequential panels |
| - Simple but complete story arc with beginning, middle, and end |
| |
| This story will unfold across the required number of scenes showing the character's journey with visual consistency and narrative coherence. |
| """ |
| |
| return enhanced.strip() |
| @log_execution |
|
|
| def _simplified_enhancement(self, user_description): |
| """ |
| Simplified enhancement fallback when the main enhancement fails. |
| |
| Args: |
| user_description: Original user story description |
| |
| Returns: |
| str: Simplified enhanced description focused on coherence for the required number of scenes. |
| """ |
| try: |
| simplified_prompt = f""" |
| Briefly enhance this story for an action-focused visual narrative. Keep it concise and coherent. |
| |
| Original: "{user_description}" |
| |
| Focus on: |
| - Core character appearance notes. |
| - Main setting description. |
| - Basic story flow suitable for action scenes. |
| - Visual consistency hints. |
| |
| Enhanced story (1-2 sentences): |
| """ |
| |
| response = self.model.generate_content(simplified_prompt) |
| enhanced_story = response.text.strip() |
| |
| if enhanced_story and len(enhanced_story) > 20: |
| print(f"[StoryGenerator] Used simplified enhancement successfully") |
| return enhanced_story |
| else: |
| return user_description |
| |
| except Exception as e: |
| print(f"[StoryGenerator] Simplified enhancement also failed: {e}") |
| return user_description |
| @log_execution |
| def generate_story(self, user_description, panels_per_page=9, num_pages=1): |
| """ |
| Generate a comic panel style story directly from user input. |
| |
| Args: |
| user_description: The user's story idea or prompt |
| panels_per_page: Number of panels per comic page (default is 8) |
| num_pages: Number of pages to generate (default is 1) |
| |
| Returns: |
| story_data: Structured data for the story with panels organized by pages |
| """ |
| enhanced_story = self.enhance_user_story(user_description) |
|
|
| panels_per_page = 9 |
| total_panels = panels_per_page * num_pages |
| print(f"[StoryGenerator] Generating comic story with {num_pages} pages, {panels_per_page} panels per page ({total_panels} total panels) from enhanced story...") |
| |
| query = f""" |
| You are a world-class comic book writer and visual storyteller. Your task is to create a SINGLE CONTINUOUS STORY. |
| The story will span exactly {num_pages} pages. Each page must contain exactly {panels_per_page} sequential action-focused panels (total of {total_panels} panels). |
| The final output must be a modern, digital-style comic with high quality and resolution, suitable for a 1024x1536 image size. **All {panels_per_page} panels must fit entirely within the page with clear gutters—no panel content may be cropped or cut off.** |
| Avoid any deformities, missing limbs, distorted or missing facial features, blurry visuals, or sketch styles. Ensure all panels are exactly the same size. |
| |
| STORY CONCEPT: |
| "{enhanced_story}" |
| |
| KEY REQUIREMENTS: |
| 1. **Panel Count & Style**: Strictly {panels_per_page} action scenes per page. No filler. All scenes must be dynamic and contribute to the story's momentum. |
| 2. **Visual Quality**: Generate ultra-high quality, modern digital comic art. Ensure no visual defects (deformities, missing limbs, distorted faces). All panels must be suitable for a combined 1024x1536 page layout. |
| 3. **Continuity**: |
| * Story must flow seamlessly page-to-page and panel-to-panel. |
| * Maintain consistent character appearances (detailed in a character sheet you will generate) and settings (detailed in a setting guide you will generate). |
| * Logical plot progression: actions have clear causes and effects. |
| * Show passage of time clearly (e.g., "later," "next day"). |
| 4. **Narrative Structure**: |
| * Complete arc: beginning, rising action, climax, resolution. |
| * Meaningful character development and motivations. |
| 5. **Visual Storytelling Focus**: |
| * Descriptions should emphasize actions, expressions, and settings to make the story understandable through visuals alone. |
| * Each panel description needs: camera angle, character positions, expressions, environment details, color palette, and mood. |
| * Focus on clear, dynamic action sequences. |
| |
| JSON OUTPUT STRUCTURE: |
| {{ |
| "title": "Overall Story Title", |
| "premise": "Brief story overview, themes, and setting.", |
| "characters": [ |
| {{ |
| "name": "Character Name", |
| "visual_description": "DETAILED visual description: height, build, face, hair, clothing. CRITICAL for consistency.", |
| "traits": ["Key visual trait 1", "Key visual trait 2"], |
| "background": "Brief backstory.", |
| "arc": "Character's journey/change." |
| }} |
| // ... (add more characters as needed) |
| ], |
| "settings": [ |
| {{ |
| "name": "Setting Name", |
| "description": "DETAILED visual description of the location, including key elements for consistency.", |
| "visual_elements": ["Notable visual element 1", "Notable visual trait 2"], |
| "mood": "Atmosphere of the location." |
| }} |
| // ... (add more settings as needed) |
| ], |
| "pages": [ |
| {{ |
| "page_number": 1, |
| "panels": [ // Exactly {panels_per_page} panels |
| {{ |
| "panel_number": 1, |
| "title": "Action-Oriented Panel Title", |
| "visual_description": "ACTION-FOCUSED, extremely detailed description: character actions, expressions, positions, environment, lighting, colors, camera angle. Ensure it fits 1024x1536 page context. NO FILLER.", |
| "text": "Dialogue/narration (context only, not for image)", |
| "purpose": "How this ACTION panel drives the story.", |
| "symbolism": "Any visual symbols." |
| }} |
| // ... (repeat for all {panels_per_page} panels on page 1) |
| ] |
| }} |
| // ... (repeat for all {num_pages} pages) |
| ] |
| }} |
| |
| REMEMBER: |
| - Focus on ACTION scenes. Eliminate all filler. |
| - Visuals are paramount. Descriptions must be rich and allow for image generation that tells the story without text. |
| - Adhere strictly to {panels_per_page} panels per page. |
| - Ensure top-tier digital art quality with no visual errors. |
| - All panels on a page contribute to a single 1024x1536 image. |
| """ |
| |
| self.log_prompt(query) |
| response = self.model.generate_content(query) |
| |
| try: |
| json_match = re.search(r'\{[\s\S]*\}', response.text, re.DOTALL) |
| if json_match: |
| json_str = json_match.group(0) |
| |
| json_str = self._fix_json(json_str) |
| |
| story_data = json.loads(json_str) |
| |
| story_data = self._validate_and_fix_structure(story_data, panels_per_page, num_pages) |
| |
| print(f"[StoryGenerator] Successfully generated story: {story_data.get('title', 'Untitled')}") |
| return story_data |
| else: |
| print("[StoryGenerator] No valid JSON found in response.") |
| raise ValueError("No valid JSON found in response") |
| except Exception as e: |
| print(f"Error in StoryGenerator: {e}") |
| return self._create_fallback_story(user_description, panels_per_page, num_pages) |
| @log_execution |
| def _validate_and_fix_structure(self, story_data, panels_per_page, num_pages): |
| """Validate and fix the story structure if needed.""" |
| if "title" not in story_data: |
| story_data["title"] = "Untitled Comic" |
| |
| if "premise" not in story_data: |
| story_data["premise"] = "A visual story." |
| |
| if "characters" not in story_data: |
| story_data["characters"] = [] |
| |
| for character in story_data.get("characters", []): |
| if "visual_description" not in character: |
| character["visual_description"] = "A character in the story." |
| if "traits" not in character: |
| character["traits"] = [] |
| if "background" not in character: |
| character["background"] = "Unknown background." |
| if "arc" not in character: |
| character["arc"] = "Experiences events throughout the story." |
| |
| if "settings" not in story_data: |
| story_data["settings"] = [] |
| |
| for setting in story_data.get("settings", []): |
| if "description" not in setting: |
| setting["description"] = "A location in the story." |
| if "visual_elements" not in setting: |
| setting["visual_elements"] = [] |
| if "mood" not in setting: |
| setting["mood"] = "Neutral." |
| |
| if "pages" not in story_data: |
| if "panels" in story_data: |
| panels = story_data.pop("panels") |
| story_data["pages"] = [] |
| |
| for i in range(num_pages): |
| start_idx = i * panels_per_page |
| end_idx = start_idx + panels_per_page |
| page_panels = panels[start_idx:end_idx] if start_idx < len(panels) else [] |
| |
| while len(page_panels) < panels_per_page: |
| panel_num = len(page_panels) + 1 + (i * panels_per_page) |
| page_panels.append({ |
| "panel_number": panel_num, |
| "title": f"Panel {panel_num}", |
| "visual_description": "A placeholder panel", |
| "text": "", |
| "purpose": "Continuation of the story", |
| "symbolism": "" |
| }) |
| |
| story_data["pages"].append({ |
| "page_number": i + 1, |
| "panels": page_panels |
| }) |
| else: |
| story_data["pages"] = [] |
| for i in range(num_pages): |
| page_panels = [] |
| for j in range(panels_per_page): |
| panel_num = j + 1 + (i * panels_per_page) |
| page_panels.append({ |
| "panel_number": panel_num, |
| "title": f"Panel {panel_num}", |
| "visual_description": "A placeholder panel", |
| "text": "", |
| "purpose": "Continuation of the story", |
| "symbolism": "" |
| }) |
| |
| story_data["pages"].append({ |
| "page_number": i + 1, |
| "panels": page_panels |
| }) |
| |
| for i in range(len(story_data["pages"]) - 1): |
| current_page = story_data["pages"][i] |
| next_page = story_data["pages"][i + 1] |
| |
| if "panels" in current_page and "panels" in next_page and current_page["panels"] and next_page["panels"]: |
| last_panel = current_page["panels"][-1] |
| first_panel = next_page["panels"][0] |
| |
| last_panel_desc = last_panel.get("visual_description", "") |
| last_panel_action = last_panel.get("text", "") |
| |
| continuity_note = f"Continues directly from page {current_page.get('page_number', i+1)}, panel {last_panel.get('panel_number', len(current_page['panels']))}: {last_panel_desc[:100]}..." |
| |
| first_panel["continuity_note"] = continuity_note |
| |
| if "visual_description" in first_panel: |
| if not first_panel["visual_description"].startswith("CONTINUING DIRECTLY"): |
| first_panel["visual_description"] = "CONTINUING DIRECTLY from previous page: " + first_panel["visual_description"] |
| |
| for i, page in enumerate(story_data["pages"]): |
| if "page_number" not in page: |
| page["page_number"] = i + 1 |
| |
| if "panels" not in page: |
| page["panels"] = [] |
| |
| if len(page["panels"]) > panels_per_page: |
| page["panels"] = page["panels"][:panels_per_page] |
| |
| while len(page["panels"]) < panels_per_page: |
| panel_num = len(page["panels"]) + 1 + (i * panels_per_page) |
| |
| context_desc = "" |
| if page["panels"]: |
| prev_panel = page["panels"][-1] |
| prev_desc = prev_panel.get("visual_description", "") |
| context_desc = f"Continuing from previous panel: {prev_desc[:50]}... " |
| |
| page["panels"].append({ |
| "panel_number": panel_num, |
| "title": f"Panel {panel_num}", |
| "visual_description": f"{context_desc}A scene related to the story, moving the narrative forward.", |
| "text": "", |
| "purpose": "Continuation of the story progression", |
| "symbolism": "" |
| }) |
| |
| for j, panel in enumerate(page["panels"]): |
| panel_num = j + 1 + (i * panels_per_page) |
| |
| if "panel_number" not in panel: |
| panel["panel_number"] = panel_num |
| |
| if "title" not in panel or not panel["title"]: |
| panel["title"] = f"Panel {panel_num}" |
| |
| if "visual_description" not in panel or not panel["visual_description"]: |
| context_desc = "" |
| if j > 0: |
| prev_panel = page["panels"][j-1] |
| prev_desc = prev_panel.get("visual_description", "") |
| context_desc = f"Following from previous panel: {prev_desc[:50]}... " |
| |
| panel["visual_description"] = f"{context_desc}A scene that advances the story narrative." |
| |
| if "text" not in panel: |
| panel["text"] = "" |
| |
| if "purpose" not in panel: |
| panel["purpose"] = "Advancing the story progression" |
| |
| if "symbolism" not in panel: |
| panel["symbolism"] = "" |
| |
| while len(story_data["pages"]) < num_pages: |
| page_num = len(story_data["pages"]) + 1 |
| page_panels = [] |
| |
| context_from_prev_page = "" |
| if story_data["pages"]: |
| prev_page = story_data["pages"][-1] |
| if prev_page.get("panels"): |
| last_panel = prev_page["panels"][-1] |
| last_desc = last_panel.get("visual_description", "") |
| context_from_prev_page = f"Continuing directly from the previous page: {last_desc[:100]}... " |
| |
| for j in range(panels_per_page): |
| panel_num = j + 1 + ((page_num - 1) * panels_per_page) |
| |
| panel_desc = "A scene that advances the story narrative." |
| if j == 0 and context_from_prev_page: |
| panel_desc = context_from_prev_page + panel_desc |
| elif j > 0 and page_panels: |
| prev_panel = page_panels[j-1] |
| prev_desc = prev_panel.get("visual_description", "") |
| panel_desc = f"Following from previous panel: {prev_desc[:50]}... " + panel_desc |
| |
| page_panels.append({ |
| "panel_number": panel_num, |
| "title": f"Panel {panel_num}", |
| "visual_description": panel_desc, |
| "text": "", |
| "purpose": "Advancing the story progression", |
| "symbolism": "" |
| }) |
| |
| story_data["pages"].append({ |
| "page_number": page_num, |
| "panels": page_panels |
| }) |
| |
| return story_data |
| @log_execution |
| def _create_fallback_story(self, user_description, panels_per_page, num_pages): |
| """Create a basic fallback story structure if generation fails.""" |
| pages = [] |
| |
| for i in range(num_pages): |
| page_panels = [] |
| for j in range(panels_per_page): |
| panel_num = j + 1 + (i * panels_per_page) |
| page_panels.append({ |
| "panel_number": panel_num, |
| "title": f"Panel {panel_num}", |
| "visual_description": f"A scene related to {user_description[:30]}...", |
| "text": f"Text for panel {panel_num}", |
| "purpose": f"Part of the story progression", |
| "symbolism": "" |
| }) |
| |
| pages.append({ |
| "page_number": i + 1, |
| "panels": page_panels |
| }) |
| |
| return { |
| "title": f"A Story About {user_description[:30]}...", |
| "premise": f"A comic story about {user_description[:50]}...", |
| "pages": pages |
| } |
| @log_execution |
|
|
| def _fix_json(self, json_str): |
| """Attempt to fix common JSON issues from LLM responses.""" |
| json_str = re.sub(r'//.*?', '', json_str) |
| json_str = re.sub(r'/\*[\s\S]*?\*/', '', json_str, flags=re.DOTALL) |
|
|
| json_str = re.sub(r'([{, ]\s*)([a-zA-Z_][a-zA-Z0-9_]*)(\s*:)', r'\1"\2"\3', json_str) |
|
|
| json_str = re.sub(r',(\s*[}\\]])', r'\1', json_str) |
| return json_str |
| @log_execution |
|
|
| def generate_panel_image_prompt(self, panel_data, style=None): |
| """Generate a prompt for image generation from panel data.""" |
| style_text = f" in {style} style" if style else "" |
| |
| prompt = f"Create a comic book panel{style_text} showing: {panel_data['visual_description']}. " |
| if 'text' in panel_data and panel_data['text']: |
| prompt += f"The panel includes the dialogue: '{panel_data['text']}'. " |
| return prompt |