Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler | |
| from PIL import Image | |
| import io | |
| import requests | |
| import os | |
| from datetime import datetime | |
| import re | |
| import tempfile | |
| import time | |
| import base64 | |
| import json | |
| from typing import Dict, List, Tuple | |
| # External OCI API URL | |
| OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space" | |
| # Try to import your existing OCI connector for direct access | |
| try: | |
| from app import oci_connector | |
| DIRECT_OCI_ACCESS = True | |
| print("✅ Direct OCI access available - using existing OCI connector") | |
| except ImportError: | |
| DIRECT_OCI_ACCESS = False | |
| print("⚠️ Direct OCI access not available - using external API endpoint") | |
| # HIGH-QUALITY MODEL SELECTION | |
| MODEL_CHOICES = { | |
| "dreamshaper-8": "lykon/dreamshaper-8", | |
| "realistic-vision": "SG161222/Realistic_Vision_V5.1", | |
| "anything-v5": "andite/anything-v5.0", | |
| "openjourney": "prompthero/openjourney", | |
| "sd-2.1": "stabilityai/stable-diffusion-2-1", | |
| } | |
| # Story tracking for sequence numbering | |
| story_registry: Dict[str, int] = {} # {story_title: current_sequence} | |
| # Initialize the HIGH-QUALITY Stable Diffusion model | |
| def load_model(model_name="dreamshaper-8"): | |
| """Load and return a high-quality Stable Diffusion model""" | |
| print(f"🔄 Loading {model_name} model...") | |
| try: | |
| model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8") | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| model_id, | |
| torch_dtype=torch.float32, | |
| safety_checker=None, | |
| requires_safety_checker=False | |
| ) | |
| pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) | |
| pipe = pipe.to("cpu") | |
| print(f"✅ {model_name} loaded successfully!") | |
| return pipe | |
| except Exception as e: | |
| print(f"❌ Model loading failed: {e}") | |
| return StableDiffusionPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| torch_dtype=torch.float32, | |
| safety_checker=None, | |
| requires_safety_checker=False | |
| ).to("cpu") | |
| # Load the model once at startup | |
| pipe = load_model() | |
| # PROFESSIONAL PROMPT ENGINEERING | |
| def enhance_prompt(prompt, style="childrens_book"): | |
| """Transform basic prompts into professional-grade prompts""" | |
| style_templates = { | |
| "childrens_book": [ | |
| "masterpiece, best quality, 4K, ultra detailed, children's book illustration", | |
| "watercolor painting, whimsical, cute, charming, storybook style", | |
| "vibrant colors, soft lighting, magical, enchanting, dreamlike", | |
| "Pixar style, Disney animation, high detail, professional artwork" | |
| ], | |
| "realistic": [ | |
| "photorealistic, 8K, ultra detailed, professional photography", | |
| "sharp focus, studio lighting, high resolution, intricate details", | |
| "realistic textures, natural lighting, cinematic quality" | |
| ], | |
| "fantasy": [ | |
| "epic fantasy art, digital painting, concept art, trending on artstation", | |
| "magical, mystical, ethereal, otherworldly, fantasy illustration", | |
| "dynamic composition, dramatic lighting, highly detailed" | |
| ], | |
| "anime": [ | |
| "anime style, Japanese animation, high quality, detailed artwork", | |
| "beautiful anime illustration, vibrant colors, clean lines", | |
| "studio ghibli style, makoto shinkai, professional anime art" | |
| ] | |
| } | |
| templates = style_templates.get(style, style_templates["childrens_book"]) | |
| style_prompt = templates[0] | |
| enhanced = f"{style_prompt}, {prompt}" | |
| quality_boosters = [ | |
| "intricate details", "beautiful composition", "perfect lighting", | |
| "professional artwork", "award winning", "trending on artstation" | |
| ] | |
| import random | |
| boosters = random.sample(quality_boosters, 2) | |
| enhanced += ", " + ", ".join(boosters) | |
| negative_prompt = ( | |
| "blurry, low quality, low resolution, ugly, deformed, poorly drawn, " | |
| "bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, " | |
| "disconnected limbs, mutation, mutated, ugly, disgusting, bad art, " | |
| "beginner, amateur, distorted, watermark, signature, text, username" | |
| ) | |
| return enhanced, negative_prompt | |
| def get_next_sequence_number(story_title: str) -> int: | |
| """Get the next sequence number for a story""" | |
| # Clean the story title for use in filenames | |
| clean_title = re.sub(r'[^a-zA-Z0-9_\-]', '', story_title.strip().replace(' ', '_')) | |
| if clean_title not in story_registry: | |
| story_registry[clean_title] = 0 | |
| story_registry[clean_title] += 1 | |
| return story_registry[clean_title], clean_title | |
| def save_to_oci_via_api(image, prompt, story_title, sequence_number): | |
| """Save image with organized storybook structure""" | |
| try: | |
| # Convert image to bytes | |
| img_bytes = io.BytesIO() | |
| image.save(img_bytes, format='PNG') | |
| img_data = img_bytes.getvalue() | |
| # Get sequence number and clean title | |
| seq_num, clean_title = get_next_sequence_number(story_title) | |
| if sequence_number is not None: | |
| seq_num = sequence_number # Use provided sequence number if available | |
| # Create organized filename with sequence | |
| filename = f"page_{seq_num:03d}_{clean_title}.png" | |
| # Use the EXTERNAL API URL | |
| api_url = f"{OCI_API_BASE_URL}/api/upload" | |
| print(f"📁 Saving to: stories/{clean_title}/{filename}") | |
| # Prepare form data for API request | |
| files = { | |
| 'file': (filename, img_data, 'image/png') | |
| } | |
| data = { | |
| 'project_id': 'storybook-library', | |
| 'subfolder': f'stories/{clean_title}' | |
| } | |
| # Make the API request | |
| response = requests.post(api_url, files=files, data=data, timeout=30) | |
| if response.status_code == 200: | |
| result = response.json() | |
| if result['status'] == 'success': | |
| return f"✅ Saved: {filename} | Story: {clean_title}" | |
| else: | |
| return f"❌ API Error: {result.get('message', 'Unknown error')}" | |
| else: | |
| return f"❌ HTTP Error: {response.status_code}" | |
| except Exception as e: | |
| return f"❌ Upload failed: {str(e)}" | |
| def generate_storybook_page(prompt, story_title, sequence_number=None, model_choice="dreamshaper-8", style="childrens_book"): | |
| """Generate a storybook page with organized naming""" | |
| try: | |
| if not prompt or not prompt.strip(): | |
| return None, "❌ Please enter a scene description" | |
| if not story_title or not story_title.strip(): | |
| return None, "❌ Please enter a story title" | |
| # Reload model if different choice | |
| global pipe | |
| if model_choice != "dreamshaper-8": | |
| pipe = load_model(model_choice) | |
| # Professional prompt enhancement | |
| enhanced_prompt, negative_prompt = enhance_prompt(prompt, style) | |
| print(f"📖 Generating page for: {story_title}") | |
| print(f"🎨 Prompt: {enhanced_prompt}") | |
| # Generate high-quality image | |
| image = pipe( | |
| prompt=enhanced_prompt, | |
| negative_prompt=negative_prompt, | |
| num_inference_steps=30, | |
| guidance_scale=8.5, | |
| width=768, | |
| height=768, | |
| generator=torch.Generator(device="cpu").manual_seed(int(time.time())) | |
| ).images[0] | |
| print("✅ Page generated successfully!") | |
| # Save with organized structure | |
| save_status = save_to_oci_via_api(image, prompt, story_title, sequence_number) | |
| print(f"💾 {save_status}") | |
| return image, save_status | |
| except Exception as e: | |
| error_msg = f"❌ Generation failed: {str(e)}" | |
| print(error_msg) | |
| return None, error_msg | |
| def batch_generate_storybook(story_title, scenes_text, model_choice="dreamshaper-8", style="childrens_book"): | |
| """Generate multiple pages for a complete storybook""" | |
| if not story_title or not scenes_text: | |
| return [], "❌ Please provide story title and scenes" | |
| scenes = [scene.strip() for scene in scenes_text.split('\n') if scene.strip()] | |
| results = [] | |
| status_messages = [] | |
| for i, scene in enumerate(scenes, 1): | |
| print(f"📖 Generating page {i}/{len(scenes)} for: {story_title}") | |
| image, status = generate_storybook_page( | |
| scene, story_title, i, model_choice, style | |
| ) | |
| if image: | |
| results.append((f"Page {i}: {scene}", image)) | |
| status_messages.append(f"Page {i}: {status}") | |
| return results, "\n".join(status_messages) | |
| # Create the enhanced Gradio interface | |
| with gr.Blocks(title="Professional Storybook Generator", theme="soft") as demo: | |
| gr.Markdown("# 📚 Professional Storybook Generator") | |
| gr.Markdown("Create organized storybooks with sequential page numbering") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 📖 Story Information") | |
| story_title = gr.Textbox( | |
| label="Story Title", | |
| placeholder="Enter your story title...\nExample: The Dragon's Reading Adventure", | |
| lines=1 | |
| ) | |
| gr.Markdown("### 🎯 Quality Settings") | |
| model_choice = gr.Dropdown( | |
| label="AI Model", | |
| choices=list(MODEL_CHOICES.keys()), | |
| value="dreamshaper-8" | |
| ) | |
| style_choice = gr.Dropdown( | |
| label="Art Style", | |
| choices=["childrens_book", "realistic", "fantasy", "anime"], | |
| value="childrens_book" | |
| ) | |
| with gr.Column(scale=2): | |
| gr.Markdown("### 🎨 Single Page Generation") | |
| prompt_input = gr.Textbox( | |
| label="Page Description", | |
| placeholder="Describe this page's scene...\nExample: The dragon discovers a magical library in the forest", | |
| lines=2 | |
| ) | |
| generate_btn = gr.Button("✨ Generate Single Page", variant="primary") | |
| image_output = gr.Image(label="Generated Page", height=400, show_download_button=True) | |
| status_output = gr.Textbox(label="Status", interactive=False, lines=2) | |
| with gr.Row(): | |
| gr.Markdown("### 📚 Complete Storybook Generation") | |
| with gr.Row(): | |
| with gr.Column(): | |
| scenes_input = gr.Textbox( | |
| label="Story Scenes (One per line = One per page)", | |
| placeholder="Enter each page's scene on a separate line...\nExample:\nThe dragon finds a mysterious book\nHe learns to read with owl friend\nThey discover hidden treasure map\nFriends celebrate with magical feast", | |
| lines=6 | |
| ) | |
| batch_btn = gr.Button("📖 Generate Complete Storybook", variant="primary") | |
| with gr.Column(): | |
| batch_status = gr.Textbox(label="Generation Status", interactive=False, lines=8) | |
| batch_gallery = gr.Gallery(label="Storybook Pages", columns=2, height=600) | |
| with gr.Accordion("📁 Folder Structure Preview", open=True): | |
| gr.Markdown(""" | |
| **Your story will be organized as:** | |
| ``` | |
| storybook-library/ | |
| └── stories/ | |
| └── The_Dragons_Reading_Adventure/ | |
| ├── page_001_The_Dragons_Reading_Adventure.png | |
| ├── page_002_The_Dragons_Reading_Adventure.png | |
| ├── page_003_The_Dragons_Reading_Adventure.png | |
| └── page_004_The_Dragons_Reading_Adventure.png | |
| ``` | |
| **Perfect for:** | |
| - ✅ Easy PDF compilation later | |
| - ✅ Automatic sequencing | |
| - ✅ Organized by story title | |
| - ✅ Ready for n8n automation | |
| """) | |
| with gr.Accordion("💡 API Usage for n8n", open=False): | |
| gr.Markdown(""" | |
| **For n8n automation, use this JSON format:** | |
| ```json | |
| { | |
| "story_title": "The Dragon's Reading Adventure", | |
| "scenes": [ | |
| "The dragon finds a mysterious old book in the forest", | |
| "He meets an owl who teaches him to read", | |
| "They discover a hidden treasure map in the book", | |
| "All animal friends celebrate with a magical feast" | |
| ], | |
| "model_choice": "dreamshaper-8", | |
| "style": "childrens_book" | |
| } | |
| ``` | |
| """) | |
| # Connect buttons to functions | |
| generate_btn.click( | |
| fn=generate_storybook_page, | |
| inputs=[prompt_input, story_title, gr.Number(visible=False), model_choice, style_choice], | |
| outputs=[image_output, status_output] | |
| ) | |
| batch_btn.click( | |
| fn=batch_generate_storybook, | |
| inputs=[story_title, scenes_input, model_choice, style_choice], | |
| outputs=[batch_gallery, batch_status] | |
| ) | |
| def get_app(): | |
| return demo | |
| if __name__ == "__main__": | |
| print("🚀 Starting Professional Storybook Generator...") | |
| print("📁 Feature: Organized story sequencing enabled") | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |