import os import PIL from io import BytesIO from google import genai from dotenv import load_dotenv import pathlib from utils import ( get_scene_splitting_prompt, get_panel_prompt, SCENE_BREAK_DELIMITER, get_regeneration_prompt ) import time from PIL import Image from reportlab.pdfgen import canvas from reportlab.lib.pagesizes import A4 import shutil import uuid import hashlib from typing import Dict, Optional # Load environment variables load_dotenv() class UserSession: """Represents a user session with isolated data and API key.""" def __init__(self, session_id: str, api_key: str): self.session_id = session_id self.api_key_hash = self._hash_api_key(api_key) # Store only hash for security self._api_key = api_key # Keep encrypted or in memory only self.created_at = time.time() self.last_activity = time.time() # Session-specific directories self.base_dir = os.path.join("data", "sessions", session_id) self.output_dir = os.path.join(self.base_dir, "output") self.user_templates_dir = os.path.join(self.base_dir, "templates") self.user_references_dir = os.path.join(self.base_dir, "references") # Create session directories os.makedirs(self.output_dir, exist_ok=True) os.makedirs(self.user_templates_dir, exist_ok=True) os.makedirs(self.user_references_dir, exist_ok=True) # Session generation data self.current_generation = { 'scenes': [], 'generated_images': [], 'chat': None, 'user_preferences': {}, 'current_template_path': os.getenv("TEMPLATE_PATH", "data/templates/template.png") } def _hash_api_key(self, api_key: str) -> str: """Create a hash of the API key for identification (not storage).""" return hashlib.sha256(api_key.encode()).hexdigest()[:16] def get_api_key(self) -> str: """Get the API key (should be implemented with proper encryption in production).""" return self._api_key def update_activity(self): """Update last activity timestamp.""" self.last_activity = time.time() def is_expired(self, timeout_hours: int = 24) -> bool: """Check if session has expired.""" return (time.time() - self.last_activity) > (timeout_hours * 3600) def cleanup(self): """Clean up session files and data.""" try: if os.path.exists(self.base_dir): shutil.rmtree(self.base_dir) print(f"Cleaned up session: {self.session_id}") except Exception as e: print(f"Error cleaning up session {self.session_id}: {e}") class SessionManager: """Manages user sessions and provides isolation between users.""" def __init__(self): self.sessions: Dict[str, UserSession] = {} self.session_timeout_hours = 4 def create_session(self, api_key: str) -> str: """Create a new session for a user with their API key.""" session_id = str(uuid.uuid4()) self.sessions[session_id] = UserSession(session_id, api_key) self._cleanup_expired_sessions() return session_id def get_session(self, session_id: str) -> Optional[UserSession]: """Get a session by ID, return None if not found or expired.""" if session_id not in self.sessions: return None session = self.sessions[session_id] if session.is_expired(self.session_timeout_hours): self._remove_session(session_id) return None session.update_activity() return session def _remove_session(self, session_id: str): """Remove and cleanup a session.""" if session_id in self.sessions: self.sessions[session_id].cleanup() del self.sessions[session_id] def _cleanup_expired_sessions(self): """Clean up all expired sessions.""" expired_sessions = [ sid for sid, session in self.sessions.items() if session.is_expired(self.session_timeout_hours) ] for session_id in expired_sessions: self._remove_session(session_id) class MangaGenerator: def __init__(self, session: UserSession): """Initialize the manga generator for a specific user session.""" self.session = session self.api_key = session.get_api_key() # Use session-specific directories self.template_path = os.getenv("TEMPLATE_PATH", "data/templates/template.png") self.output_dir = session.output_dir self.user_templates_dir = session.user_templates_dir self.user_references_dir = session.user_references_dir self.image_gen_model_name = os.getenv("IMAGE_MODEL_NAME", "gemini-2.5-flash-image-preview") self.scene_gen_model_name = os.getenv("SCENE_MODEL_NAME", "gemini-2.0-flash") # Initialize clients with session's API key try: self.image_gen_client = genai.Client(api_key=self.api_key) self.scene_gen_client = genai.Client(api_key=self.api_key) except Exception as e: raise ValueError(f"Invalid API key or client initialization failed: {e}") @property def current_generation(self): """Access current generation data from session.""" return self.session.current_generation def save_user_template(self, uploaded_file): """Save user uploaded template and return the path.""" if uploaded_file is None: return None try: # Generate unique filename with timestamp timestamp = int(time.time()) filename = f"user_template_{timestamp}.png" template_path = os.path.join(self.user_templates_dir, filename) # Handle different file input types if hasattr(uploaded_file, 'name'): # Gradio file object shutil.copy2(uploaded_file.name, template_path) else: if isinstance(uploaded_file, str): shutil.copy2(uploaded_file, template_path) elif hasattr(uploaded_file, 'save'): # PIL Image uploaded_file.save(template_path) # Verify the template was saved and is a valid image test_image = PIL.Image.open(template_path) test_image.verify() print(f"User template saved to: {template_path}") return template_path except Exception as e: print(f"Error saving user template: {e}") return None def save_user_reference_image(self, uploaded_file): """Save user uploaded reference image and return the path.""" if uploaded_file is None: return None try: timestamp = int(time.time()) filename = f"user_reference_{timestamp}.png" reference_path = os.path.join(self.user_references_dir, filename) if hasattr(uploaded_file, 'name'): # Gradio file object shutil.copy2(uploaded_file.name, reference_path) else: if isinstance(uploaded_file, str): shutil.copy2(uploaded_file, reference_path) elif hasattr(uploaded_file, 'save'): # PIL Image uploaded_file.save(reference_path) test_image = PIL.Image.open(reference_path) test_image.verify() print(f"User reference image saved to: {reference_path}") return reference_path except Exception as e: print(f"Error saving user reference image: {e}") return None def set_template_for_generation(self, template_path): """Set the template to use for the current generation session.""" if template_path and os.path.exists(template_path): self.current_generation['current_template_path'] = template_path return True return False def get_current_template_path(self): """Get the current template path being used.""" return self.current_generation.get('current_template_path', self.template_path) def read_story(self, file_path): """Read story text from file.""" with open(file_path, "r") as f: return f.read() def split_into_scenes(self, story_text: str, n_scenes: int): """Split story into visual scenes with descriptions.""" prompt = get_scene_splitting_prompt(story_text, n_scenes) response = self.scene_gen_client.models.generate_content( model=self.scene_gen_model_name, contents=[prompt] ) full_response_text = "" for part in response.candidates[0].content.parts: if part.text: full_response_text += part.text scenes = [scene.strip() for scene in full_response_text.split(SCENE_BREAK_DELIMITER)] scenes = [scene for scene in scenes if scene] return scenes[:n_scenes] def save_image(self, response, path): """Save the generated image from response.""" time.sleep(3) for part in response.parts: if image := part.as_image(): image.save(path) return image return None def generate_image_for_scene(self, scene_description: str, output_path: str): """Generate image for a single scene.""" current_template = self.get_current_template_path() response = self.image_gen_client.models.generate_content( model=self.image_gen_model_name, contents=[ scene_description, PIL.Image.open(current_template) ] ) saved_image = self.save_image(response, output_path) return response, saved_image def generate_image_with_chat(self, scene_description: str, output_path: str, chat): """Generate image using chat context for consistency.""" current_template = self.get_current_template_path() response = chat.send_message([ scene_description, PIL.Image.open(current_template) ]) saved_image = self.save_image(response, output_path) return response, saved_image def generate_image_with_chat_and_reference(self, scene_description: str, output_path: str, chat, reference_image_path=None): """Generate image using chat context with optional reference image.""" current_template = self.get_current_template_path() content = [scene_description, PIL.Image.open(current_template)] if reference_image_path and os.path.exists(reference_image_path): content.append(PIL.Image.open(reference_image_path)) print(f"Using reference image: {reference_image_path}") response = chat.send_message(content) saved_image = self.save_image(response, output_path) return response, saved_image def regenerate_specific_panel(self, panel_index: int, modification_request: str, reference_image=None): """Regenerate a specific panel with modifications and optional reference image.""" if not self.current_generation['scenes'] or not self.current_generation['chat']: raise ValueError("No active generation session. Please generate manga first.") if panel_index >= len(self.current_generation['scenes']): raise ValueError(f"Panel index {panel_index} is out of range.") original_scene = self.current_generation['scenes'][panel_index] reference_image_path = None if reference_image is not None: reference_image_path = self.save_user_reference_image(reference_image) if reference_image_path: modification_request += "\n\nIMPORTANT: Use the provided reference image as visual guidance for style, composition, or specific elements while maintaining the story's integrity." user_preferences = self.current_generation.get('user_preferences', {}) modified_prompt = get_regeneration_prompt( original_scene, modification_request, is_first_panel=(panel_index == 0), user_preferences=user_preferences ) current_version = self.current_generation['generated_images'][panel_index].get('version', 1) output_path = os.path.join(self.output_dir, f"scene{panel_index+1}_v{current_version + 1}.png") response, saved_image = self.generate_image_with_chat_and_reference( modified_prompt, output_path, self.current_generation['chat'], reference_image_path ) return output_path, saved_image def replace_panel(self, panel_index: int, new_image_path: str, new_image: PIL.Image): """Replace a panel in the current generation.""" if not self.current_generation['generated_images']: raise ValueError("No active generation session.") if panel_index >= len(self.current_generation['generated_images']): raise ValueError(f"Panel index {panel_index} is out of range.") current_version = self.current_generation['generated_images'][panel_index].get('version', 1) self.current_generation['generated_images'][panel_index].update({ 'image_path': new_image_path, 'image': new_image, 'version': current_version + 1 }) def get_current_gallery_paths(self): """Get current image paths for gallery display.""" if not self.current_generation['generated_images']: return [] return [img['image_path'] for img in self.current_generation['generated_images']] def generate_manga_from_story(self, story_text: str, n_scenes: int = 5, user_preferences: dict = None, user_template=None): """Generate complete manga from story text with user preferences and optional custom template.""" if user_template is not None: template_path = self.save_user_template(user_template) if template_path: self.set_template_for_generation(template_path) print(f"Using user template: {template_path}") else: print("Failed to save user template, using default template") if user_preferences is None: user_preferences = {} scenes = self.split_into_scenes(story_text, n_scenes) chat = self.image_gen_client.chats.create(model=self.image_gen_model_name) generated_images = [] responses = [] for i, scene in enumerate(scenes): scene_description = get_panel_prompt(scene, is_first_panel=(i == 0), user_preferences=user_preferences) output_path = os.path.join(self.output_dir, f"scene{i+1}.png") response, saved_image = self.generate_image_with_chat(scene_description, output_path, chat) responses.append(response) generated_images.append({ 'scene_number': i + 1, 'scene_text': scene, 'image_path': output_path, 'image': saved_image, 'version': 1 }) print(f"Generated scene {i+1}") self.current_generation.update({ 'scenes': scenes, 'generated_images': generated_images, 'chat': chat, 'user_preferences': user_preferences }) return generated_images, scenes def generate_manga_from_file(self, story_file_path: str, n_scenes: int = 5, user_preferences: dict = None, user_template=None): """Generate manga from story file with user preferences and optional custom template.""" story_text = self.read_story(story_file_path) return self.generate_manga_from_story(story_text, n_scenes, user_preferences, user_template) def get_current_panels(self): """Get current panel information for the interface.""" if not self.current_generation['generated_images']: return [] return [(i+1, img['image_path']) for i, img in enumerate(self.current_generation['generated_images'])] def create_manga_pdf(self, output_filename=None): """Create a PDF file from all current manga panels.""" if not self.current_generation['generated_images']: raise ValueError("No manga panels to export. Please generate manga first.") if output_filename is None: output_filename = os.path.join(self.output_dir, "manga_complete.pdf") c = canvas.Canvas(output_filename, pagesize=A4) page_width, page_height = A4 c.setFont("Helvetica-Bold", 24) title_text = "Generated Manga" title_width = c.stringWidth(title_text, "Helvetica-Bold", 24) c.drawString((page_width - title_width) / 2, page_height - 100, title_text) c.setFont("Helvetica", 12) subtitle_text = f"Total Panels: {len(self.current_generation['generated_images'])}" subtitle_width = c.stringWidth(subtitle_text, "Helvetica", 12) c.drawString((page_width - subtitle_width) / 2, page_height - 130, subtitle_text) current_template = self.get_current_template_path() if current_template != self.template_path: template_info = "Custom Template Used" template_width = c.stringWidth(template_info, "Helvetica", 12) c.drawString((page_width - template_width) / 2, page_height - 150, template_info) c.showPage() for i, panel_data in enumerate(self.current_generation['generated_images']): image_path = panel_data['image_path'] if os.path.exists(image_path): img = Image.open(image_path) img_width, img_height = img.size aspect_ratio = img_width / img_height max_width = page_width - 100 max_height = page_height - 150 if aspect_ratio > 1: new_width = min(max_width, img_width) new_height = new_width / aspect_ratio else: new_height = min(max_height, img_height) new_width = new_height * aspect_ratio x = (page_width - new_width) / 2 y = (page_height - new_height) / 2 c.drawImage(image_path, x, y, width=new_width, height=new_height) c.setFont("Helvetica", 10) c.drawString(50, page_height - 30, f"Panel {i + 1}") c.showPage() c.save() print(f"PDF saved to: {output_filename}") return output_filename # Global session manager _session_manager = SessionManager() def get_session_manager() -> SessionManager: """Get the global session manager.""" return _session_manager def get_generator_for_session(session_id: str) -> Optional[MangaGenerator]: """Get a manga generator for a specific session.""" session = _session_manager.get_session(session_id) if session is None: return None return MangaGenerator(session) def create_user_session(api_key: str) -> str: """Create a new user session with API key.""" if not api_key or not api_key.strip(): raise ValueError("API key is required") return _session_manager.create_session(api_key.strip()) # Interface functions with session support def generate_manga_interface(session_id: str, story_text: str, num_scenes: int = 5, art_style: str = None, mood: str = None, color_palette: str = None, character_style: str = None, line_style: str = None, composition: str = None, additional_notes: str = "", user_template=None): """Interface function for Gradio - generates manga from text input with session support.""" try: generator = get_generator_for_session(session_id) if generator is None: return [], "Session expired or invalid. Please refresh and enter your API key again." user_preferences = {} if art_style and art_style != "None": user_preferences['art_style'] = art_style if mood and mood != "None": user_preferences['mood'] = mood if color_palette and color_palette != "None": user_preferences['color_palette'] = color_palette if character_style and character_style != "None": user_preferences['character_style'] = character_style if line_style and line_style != "None": user_preferences['line_style'] = line_style if composition and composition != "None": user_preferences['composition'] = composition if additional_notes.strip(): user_preferences['additional_notes'] = additional_notes.strip() generated_images, scenes = generator.generate_manga_from_story(story_text, num_scenes, user_preferences, user_template) image_paths = [img['image_path'] for img in generated_images] scene_descriptions = [img['scene_text'] for img in generated_images] return image_paths, "\n\n".join([f"Scene {i+1}: {scene}" for i, scene in enumerate(scene_descriptions)]) except Exception as e: print(f"Error generating manga: {e}") return [], f"Error: {str(e)}" def generate_manga_from_file_interface(session_id: str, story_file, num_scenes: int = 5, art_style: str = None, mood: str = None, color_palette: str = None, character_style: str = None, line_style: str = None, composition: str = None, additional_notes: str = "", user_template=None): """Interface function for Gradio - generates manga from uploaded file with session support.""" try: generator = get_generator_for_session(session_id) if generator is None: return [], "Session expired or invalid. Please refresh and enter your API key again." if hasattr(story_file, 'name'): # Gradio file object with open(story_file.name, 'r') as f: story_text = f.read() else: story_text = str(story_file) user_preferences = {} if art_style and art_style != "None": user_preferences['art_style'] = art_style if mood and mood != "None": user_preferences['mood'] = mood if color_palette and color_palette != "None": user_preferences['color_palette'] = color_palette if character_style and character_style != "None": user_preferences['character_style'] = character_style if line_style and line_style != "None": user_preferences['line_style'] = line_style if composition and composition != "None": user_preferences['composition'] = composition if additional_notes.strip(): user_preferences['additional_notes'] = additional_notes.strip() generated_images, scenes = generator.generate_manga_from_story(story_text, num_scenes, user_preferences, user_template) image_paths = [img['image_path'] for img in generated_images] scene_descriptions = [img['scene_text'] for img in generated_images] return image_paths, "\n\n".join([f"Scene {i+1}: {scene}" for i, scene in enumerate(scene_descriptions)]) except Exception as e: print(f"Error generating manga: {e}") return [], f"Error: {str(e)}" def regenerate_and_replace_interface(session_id: str, panel_number: int, modification_request: str, replace_original: bool, reference_image=None): """Interface function for Gradio - regenerate panel with session support.""" try: generator = get_generator_for_session(session_id) if generator is None: return "Session expired or invalid. Please refresh and enter your API key again.", None, [] if not modification_request.strip(): return "Please provide modification instructions.", None, [] panel_index = panel_number - 1 new_image_path, saved_image = generator.regenerate_specific_panel(panel_index, modification_request, reference_image) updated_gallery = [] if replace_original and generator.current_generation['generated_images']: generator.replace_panel(panel_index, new_image_path, saved_image) updated_gallery = generator.get_current_gallery_paths() status_message = f"Panel {panel_number} regenerated and replaced successfully!" else: status_message = f"Panel {panel_number} regenerated successfully! (Original preserved)" return status_message, new_image_path, updated_gallery except Exception as e: return f"Error regenerating panel: {e}", None, [] def create_pdf_interface(session_id: str): """Interface function for Gradio - create PDF with session support.""" try: generator = get_generator_for_session(session_id) if generator is None: return "Session expired or invalid. Please refresh and enter your API key again.", None if not generator.current_generation['generated_images']: return "No manga panels to export. Please generate manga first.", None pdf_path = generator.create_manga_pdf() message = f"PDF created successfully! ({len(generator.current_generation['generated_images'])} panels)" return message, pdf_path except Exception as e: return f"Error creating PDF: {e}", None def get_current_panels(session_id: str): """Get current panel information for the interface with session support.""" generator = get_generator_for_session(session_id) if generator is None: return [] return generator.get_current_panels()