Update app.py

app.py

@@ -98,7 +98,7 @@ class JobStatus(str, Enum):
     FAILED = "failed"
 
 # =============================================
-# ENHANCED: Style and Size Configuration
+# Style and Size Configuration
 # =============================================
 
 # Available image styles with their prompts
@@ -185,7 +185,7 @@ STYLE_CONFIGS = {
     }
 }
 
-# Size presets
+# Size presets with valid dimensions (divisible by 8)
 SIZE_PRESETS = {
     "square": {"width": 768, "height": 768, "description": "1:1 Square"},
     "portrait": {"width": 768, "height": 1024, "description": "3:4 Portrait"},
@@ -194,8 +194,8 @@ SIZE_PRESETS = {
     "ultrawide": {"width": 1920, "height": 768, "description": "2.5:1 Ultrawide"},
     "instagram": {"width": 1080, "height": 1080, "description": "1:1 Instagram"},
     "story": {"width": 1080, "height": 1920, "description": "9:16 Story"},
-    "facebook": {"width": 1200, "height": 630, "description": "1.91:1 Facebook"},
-    "twitter": {"width": 1600, "height": 900, "description": "16:9 Twitter"},
+    "facebook": {"width": 1200, "height": 632, "description": "1.91:1 Facebook"},
+    "twitter": {"width": 1600, "height": 896, "description": "16:9 Twitter"},
     "custom": {"width": 768, "height": 1024, "description": "Custom Size"}
 }
 
@@ -214,9 +214,9 @@ class StorybookRequest(BaseModel):
     characters: List[CharacterDescription] = []
     model_choice: str = "dreamshaper-8"
     style: str = "childrens_book"
-    image_size: str = "portrait"  # New: size preset
-    custom_width: Optional[int] = None  # New: custom width
-    custom_height: Optional[int] = None  # New: custom height
+    image_size: str = "portrait"
+    custom_width: Optional[int] = None
+    custom_height: Optional[int] = None
     callback_url: Optional[str] = None
     consistency_seed: Optional[int] = None
     project_id: Optional[str] = None
@@ -230,7 +230,7 @@ class InfographicRequest(BaseModel):
     title: str
     items: List[InfographicItem]
     style: str = "food_photography"
-    image_size: str = "story"  # New: size preset for infographic
+    image_size: str = "story"
     custom_width: Optional[int] = None
     custom_height: Optional[int] = None
     project_id: Optional[str] = None
@@ -275,41 +275,51 @@ current_pipe = None
 model_lock = threading.Lock()
 
 # =============================================
-# HELPER FUNCTIONS
+# HELPER FUNCTIONS WITH DIMENSION VALIDATION
 # =============================================
 
+def validate_dimensions(width, height):
+    """Ensure dimensions are compatible with Stable Diffusion (divisible by 8)"""
+    # SD models require dimensions divisible by 8
+    width = (width // 8) * 8
+    height = (height // 8) * 8
+    
+    # Minimum and maximum constraints
+    width = max(256, min(1024, width))
+    height = max(256, min(1024, height))
+    
+    return width, height
+
 def get_image_dimensions(size_preset, custom_width=None, custom_height=None):
-    """Get image dimensions based on preset or custom values"""
+    """Get image dimensions based on preset or custom values with validation"""
     if size_preset == "custom" and custom_width and custom_height:
-        return custom_width, custom_height
+        width, height = custom_width, custom_height
     elif size_preset in SIZE_PRESETS:
         preset = SIZE_PRESETS[size_preset]
-        return preset["width"], preset["height"]
+        width, height = preset["width"], preset["height"]
     else:
         # Default to portrait
-        return 768, 1024
+        width, height = 768, 1024
+    
+    # Validate dimensions
+    width, height = validate_dimensions(width, height)
+    return width, height
 
 def get_style_config(style_name):
     """Get style configuration or return default"""
     if style_name in STYLE_CONFIGS:
         return STYLE_CONFIGS[style_name]
-    # Default to children's book style
     return STYLE_CONFIGS["childrens_book"]
 
 def enhance_prompt_with_style(scene_visual, style="childrens_book", custom_style_prompt=None):
     """Enhanced prompt with style configuration"""
-    
-    # Get style configuration
     if custom_style_prompt:
         style_prefix = custom_style_prompt
     else:
         style_config = get_style_config(style)
         style_prefix = style_config["prompt_prefix"]
     
-    # Combine with user prompt
     enhanced_prompt = f"{style_prefix}, {scene_visual}"
-    
-    # Get negative prompt from style config
     negative_prompt = style_config.get("negative", "blurry, low quality, bad anatomy, deformed, ugly")
     
     return enhanced_prompt, negative_prompt
@@ -318,6 +328,9 @@ def generate_image_with_settings(prompt, model_choice, style, scene_number, cons
                                  width=768, height=1024, custom_style_prompt=None):
     """Generate image with customizable style and dimensions"""
     
+    # Validate dimensions before generation
+    width, height = validate_dimensions(width, height)
+    
     # Get enhanced prompt with style
     enhanced_prompt, negative_prompt = enhance_prompt_with_style(prompt, style, custom_style_prompt)
     
@@ -330,6 +343,12 @@ def generate_image_with_settings(prompt, model_choice, style, scene_number, cons
     try:
         pipe = load_model(model_choice)
         
+        if pipe is None:
+            raise Exception("Model pipeline not loaded")
+        
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(scene_seed)
+        
         image = pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
@@ -337,18 +356,20 @@ def generate_image_with_settings(prompt, model_choice, style, scene_number, cons
             guidance_scale=7.5,
             width=width,
             height=height,
-            generator=torch.Generator(device="cpu").manual_seed(scene_seed)
+            generator=generator
         ).images[0]
         
-        print(f"✅ Generated image - Style: {style}, Size: {width}x{height}")
+        print(f"✅ Generated image - Style: {style}, Size: {width}x{height}, Seed: {scene_seed}")
         return image
         
     except Exception as e:
         print(f"❌ Generation failed: {str(e)}")
+        print(f"   - Style: {style}")
+        print(f"   - Size: {width}x{height}")
         raise
 
 # =============================================
-# EXISTING FUNCTIONS (Keep your existing functions)
+# DATASET FUNCTIONS
 # =============================================
 
 def ensure_dataset_exists():
@@ -406,6 +427,10 @@ def upload_image_to_hf_dataset(image, project_id, page_number, prompt, style="")
         print(f"❌ Failed to upload: {e}")
         return None
 
+# =============================================
+# MEMORY MANAGEMENT
+# =============================================
+
 def get_memory_usage():
     """Get current memory usage statistics"""
     process = psutil.Process()
@@ -481,6 +506,10 @@ print("🚀 Initializing Enhanced Generator API...")
 load_model("dreamshaper-8")
 print("✅ Model loaded and ready!")
 
+# =============================================
+# LOCAL FILE MANAGEMENT
+# =============================================
+
 def save_image_to_local(image, prompt, style="test"):
     """Save image to local storage"""
     try:
@@ -548,6 +577,10 @@ def refresh_local_images():
     except Exception as e:
         return []
 
+# =============================================
+# JOB MANAGEMENT
+# =============================================
+
 def create_job(request) -> str:
     job_id = str(uuid.uuid4())
     job_storage[job_id] = {
@@ -578,7 +611,7 @@ def update_job_status(job_id: str, status: JobStatus, progress: int, message: st
     return True
 
 # =============================================
-# ENHANCED: Storybook Generation with Style & Size
+# STORYBOOK GENERATION
 # =============================================
 
 def generate_storybook_background(job_id: str):
@@ -589,14 +622,13 @@ def generate_storybook_background(job_id: str):
         
         project_id = story_request.project_id or story_request.story_title.replace(' ', '_').lower()
         
-        # Get dimensions based on size preset
         width, height = get_image_dimensions(
             story_request.image_size,
             story_request.custom_width,
             story_request.custom_height
         )
         
-        print(f"🎬 Starting storybook generation - Style: {story_request.style}, Size: {width}x{height}")
+        print(f"🎬 Starting storybook - Style: {story_request.style}, Size: {width}x{height}")
         
         update_job_status(job_id, JobStatus.PROCESSING, 5, "Starting generation...")
         
@@ -609,7 +641,6 @@ def generate_storybook_background(job_id: str):
             update_job_status(job_id, JobStatus.PROCESSING, progress, f"Generating page {i+1}/{total_scenes}")
             
             try:
-                # Generate image with style and size
                 image = generate_image_with_settings(
                     scene.visual,
                     story_request.model_choice,
@@ -659,24 +690,204 @@ def generate_storybook_background(job_id: str):
         update_job_status(job_id, JobStatus.FAILED, 0, f"Failed: {str(e)}")
 
 # =============================================
-# ENHANCED: Infographic Generation with Style & Size
+# INFOGRAPHIC LAYOUT GENERATOR
 # =============================================
 
 class InfographicLayoutGenerator:
-    # Keep your existing InfographicLayoutGenerator class here
-    # (I'm keeping it minimal for brevity - include your full class)
     def __init__(self):
         self.title_color = (60, 60, 60)
         self.text_color = (40, 40, 40)
         self.accent_color = (200, 160, 120)
     
+    def get_font(self, size, bold=False):
+        """Simple font loading"""
+        try:
+            return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", size)
+        except:
+            try:
+                return ImageFont.truetype("arial.ttf", size)
+            except:
+                return ImageFont.load_default()
+    
+    def smart_remove_background(self, image):
+        """Simple background removal"""
+        try:
+            if image.mode != 'RGBA':
+                image = image.convert('RGBA')
+            
+            data = image.getdata()
+            new_data = []
+            
+            width, height = image.size
+            
+            # Sample edges
+            edge_samples = []
+            for x in range(0, width, 10):
+                if x < len(data):
+                    edge_samples.append(data[x])
+            for x in range(0, width, 10):
+                idx = (height-1) * width + x
+                if idx < len(data):
+                    edge_samples.append(data[idx])
+            
+            if not edge_samples:
+                return image
+            
+            avg_r = sum(s[0] for s in edge_samples) // len(edge_samples)
+            avg_g = sum(s[1] for s in edge_samples) // len(edge_samples)
+            avg_b = sum(s[2] for s in edge_samples) // len(edge_samples)
+            
+            threshold = 60
+            
+            for item in data:
+                distance = ((item[0] - avg_r) ** 2 + 
+                          (item[1] - avg_g) ** 2 + 
+                          (item[2] - avg_b) ** 2) ** 0.5
+                
+                if distance < threshold:
+                    new_data.append((255, 255, 255, 0))
+                else:
+                    new_data.append(item)
+            
+            image.putdata(new_data)
+            return image
+            
+        except Exception as e:
+            return image
+    
+    def create_responsive_layout(self, num_items):
+        """Calculate responsive layout"""
+        width = 900
+        height = 1600
+        
+        if num_items <= 3:
+            icon_size = 180
+            columns = 1
+            rows = num_items
+        elif num_items <= 6:
+            icon_size = 150
+            columns = 2
+            rows = (num_items + 1) // 2
+        else:
+            icon_size = 120
+            columns = 2
+            rows = (num_items + 1) // 2
+        
+        positions = []
+        available_height = height - 300
+        start_y = 200
+        
+        for i in range(num_items):
+            col = i % columns
+            row = i // columns
+            
+            if columns == 1:
+                x = (width - icon_size) // 2
+            else:
+                if col == 0:
+                    x = width // 4 - icon_size // 2
+                else:
+                    x = 3 * width // 4 - icon_size // 2
+            
+            if rows > 1:
+                y = start_y + (row * (available_height // rows))
+            else:
+                y = start_y + (available_height - icon_size) // 2
+            
+            positions.append((x, y, icon_size))
+        
+        return positions, icon_size, columns
+    
     def create_food_infographic(self, title, items_with_prompts, generated_images, background_image=None, width=900, height=1600):
         """Create infographic with custom dimensions"""
-        # Your existing implementation here
-        canvas = Image.new('RGB', (width, height), (255, 255, 255))
-        # ... rest of your infographic code
+        if background_image:
+            bg = background_image.resize((width, height), Image.Resampling.LANCZOS)
+            canvas = bg.convert('RGB')
+        else:
+            canvas = Image.new('RGB', (width, height), (255, 255, 255))
+        
+        draw = ImageDraw.Draw(canvas)
+        num_items = len(items_with_prompts)
+        positions, icon_size, columns = self.create_responsive_layout(num_items)
+        
+        # Draw title
+        title_font = self.get_font(48, bold=True)
+        title_bbox = draw.textbbox((0, 0), title, font=title_font)
+        title_width = title_bbox[2] - title_bbox[0]
+        title_x = (width - title_width) // 2
+        title_y = 60
+        
+        draw.rounded_rectangle(
+            [title_x-20, title_y-10, title_x+title_width+20, title_y+60],
+            radius=20, fill=(250, 250, 250), outline=None
+        )
+        draw.text((title_x, title_y), title, font=title_font, fill=self.title_color)
+        
+        # Draw subtitle
+        subtitle_font = self.get_font(24)
+        subtitle = "Healthy Foods · 健康食品"
+        subtitle_bbox = draw.textbbox((0, 0), subtitle, font=subtitle_font)
+        subtitle_width = subtitle_bbox[2] - subtitle_bbox[0]
+        subtitle_x = (width - subtitle_width) // 2
+        subtitle_y = title_y + 60
+        draw.text((subtitle_x, subtitle_y), subtitle, font=subtitle_font, fill=(120, 120, 120))
+        
+        # Place items
+        for i, ((food_name, _), img) in enumerate(zip(items_with_prompts, generated_images)):
+            if i >= len(positions):
+                break
+                
+            x, y, size = positions[i]
+            
+            if img:
+                icon_img = self.smart_remove_background(img)
+                icon_img.thumbnail((size, size), Image.Resampling.LANCZOS)
+                
+                mask = Image.new('L', (size, size), 0)
+                mask_draw = ImageDraw.Draw(mask)
+                mask_draw.ellipse((0, 0, size, size), fill=255)
+                
+                icon_final = Image.new('RGBA', (size, size), (0, 0, 0, 0))
+                paste_x = (size - icon_img.width) // 2
+                paste_y = (size - icon_img.height) // 2
+                icon_final.paste(icon_img, (paste_x, paste_y), icon_img)
+                
+                circle_bg = Image.new('RGBA', (size+20, size+20), (0, 0, 0, 0))
+                circle_draw = ImageDraw.Draw(circle_bg)
+                circle_draw.ellipse([5, 5, size+15, size+15], fill=(245, 245, 245), outline=None)
+                canvas.paste(circle_bg, (x-10, y-10), circle_bg)
+                canvas.paste(icon_final, (x, y), icon_final)
+            
+            name_font = self.get_font(28 if num_items > 6 else 32, bold=True)
+            
+            if columns == 1:
+                name_x = x + (size // 2)
+                name_y = y + size + 15
+                draw.text((name_x, name_y), food_name, font=name_font, fill=self.text_color, anchor="mm")
+                number_font = self.get_font(20)
+                draw.text((name_x, name_y + 35), f"#{i+1:02d}", font=number_font, fill=(150, 150, 150), anchor="mm")
+            else:
+                name_x = x + size + 15
+                name_y = y + (size // 2) - 20
+                draw.text((name_x, name_y), food_name, font=name_font, fill=self.text_color)
+                number_font = self.get_font(18)
+                draw.text((name_x, name_y + 30), f"#{i+1:02d}", font=number_font, fill=(150, 150, 150))
+        
+        # Draw footer
+        footer_font = self.get_font(20)
+        footer_text = "·  Fresh & Healthy  ·"
+        footer_bbox = draw.textbbox((0, 0), footer_text, font=footer_font)
+        footer_width = footer_bbox[2] - footer_bbox[0]
+        footer_x = (width - footer_width) // 2
+        footer_y = height - 50
+        draw.text((footer_x, footer_y), footer_text, font=footer_font, fill=(150, 150, 150))
+        
         return canvas
 
+# =============================================
+# INFOGRAPHIC GENERATION
+# =============================================
+
 def generate_infographic_background(job_id: str):
     """Background task with style and size options"""
     try:
@@ -684,7 +895,6 @@ def generate_infographic_background(job_id: str):
         request = InfographicRequest(**job_data["request"])
         project_id = request.project_id or f"infographic_{int(time.time())}"
         
-        # Get dimensions for infographic (usually story/portrait)
         width, height = get_image_dimensions(
             request.image_size,
             request.custom_width,
@@ -715,14 +925,13 @@ def generate_infographic_background(job_id: str):
             update_job_status(job_id, JobStatus.PROCESSING, progress, f"Generating {item.name}")
             
             try:
-                # Use the enhanced generation with style
                 image = generate_image_with_settings(
                     item.prompt,
                     "dreamshaper-8",
-                    "realistic",  # Food photos use realistic style
+                    "realistic",
                     i + 1,
                     None,
-                    768, 1024  # Standard size for food images
+                    768, 1024
                 )
                 generated_images.append(image)
                 
@@ -775,7 +984,7 @@ def generate_infographic_background(job_id: str):
             "image_size": f"{width}x{height}"
         }
         
-        update_job_status(job_id, JobStatus.COMPLETED, 100, f"Completed!", result)
+        update_job_status(job_id, JobStatus.COMPLETED, 100, "Completed!", result)
         
     except Exception as e:
         update_job_status(job_id, JobStatus.FAILED, 0, f"Failed: {str(e)}")
@@ -788,7 +997,7 @@ def generate_infographic_background(job_id: str):
 async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
     """Main endpoint with style and size options"""
     try:
-        print(f"📥 Received storybook request - Style: {request.get('style', 'childrens_book')}, Size: {request.get('image_size', 'portrait')}")
+        print(f"📥 Received storybook request - Style: {request.get('style', 'childrens_book')}")
         
         if 'consistency_seed' not in request:
             request['consistency_seed'] = random.randint(1000, 9999)
@@ -803,7 +1012,6 @@ async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
         job_id = create_job(story_request)
         background_tasks.add_task(generate_storybook_background, job_id)
         
-        # Get dimensions for response
         width, height = get_image_dimensions(
             request.get('image_size', 'portrait'),
             request.get('custom_width'),
@@ -826,7 +1034,7 @@ async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
 async def generate_infographic(request: dict, background_tasks: BackgroundTasks):
     """Generate infographic with style and size options"""
     try:
-        print(f"📥 Received infographic request - Style: {request.get('style', 'food_photography')}, Size: {request.get('image_size', 'story')}")
+        print(f"📥 Received infographic request - Style: {request.get('style', 'food_photography')}")
         
         if 'project_id' not in request:
             request['project_id'] = f"infographic_{int(time.time())}"
@@ -837,7 +1045,6 @@ async def generate_infographic(request: dict, background_tasks: BackgroundTasks)
         job_id = create_job(infographic_request)
         background_tasks.add_task(generate_infographic_background, job_id)
         
-        # Get dimensions for response
         width, height = get_image_dimensions(
             request.get('image_size', 'story'),
             request.get('custom_width'),
@@ -922,13 +1129,11 @@ def test_generation(prompt, model_choice, style, size_preset, custom_width, cust
     if not prompt.strip():
         return None, "❌ Please enter a prompt"
     try:
-        # Get dimensions
         if size_preset == "custom" and custom_width and custom_height:
             width, height = custom_width, custom_height
         else:
             width, height = SIZE_PRESETS.get(size_preset, (768, 1024))
         
-        # Generate image
         image = generate_image_with_settings(
             prompt, model_choice, style, 1, None, width, height
         )