Update image_generator.py

image_generator.py

@@ -12,41 +12,47 @@ class ImageGenerator:
         self.model = None
         self.inference_steps = 20
         self.target_size = (384, 384)
+        self.aspect_ratio = "1:1"  # Default aspect ratio
         
     def load_model(self):
         """Load a lightweight image generation model"""
         if self.model is None:
             with st.spinner("Loading image generation model... This may take a moment."):
-                # Using a lightweight model for image generation
-                from diffusers import StableDiffusionPipeline
-                
-                model_id = "runwayml/stable-diffusion-v1-5"
-                
-                # Load with memory optimization settings
-                self.model = StableDiffusionPipeline.from_pretrained(
-                    model_id,
-                    torch_dtype=torch.float32,
-                    safety_checker=None,
-                    requires_safety_checker=False
-                )
-                
-                # Use CPU for inference to save memory
-                self.model = self.model.to("cpu")
-                
-                # Enable memory efficient attention if available
-                if hasattr(self.model, 'enable_attention_slicing'):
-                    self.model.enable_attention_slicing()
-                
-                # Enable memory efficient attention
-                if hasattr(self.model, 'enable_vae_slicing'):
-                    self.model.enable_vae_slicing()
-                
-                # Enable xformers memory efficient attention if available
                 try:
-                    if hasattr(self.model, 'enable_xformers_memory_efficient_attention'):
-                        self.model.enable_xformers_memory_efficient_attention()
-                except:
-                    pass
+                    # Using a lightweight model for image generation
+                    from diffusers import StableDiffusionPipeline
+                    
+                    model_id = "runwayml/stable-diffusion-v1-5"
+                    
+                    # Load with memory optimization settings
+                    self.model = StableDiffusionPipeline.from_pretrained(
+                        model_id,
+                        torch_dtype=torch.float32,
+                        safety_checker=None,
+                        requires_safety_checker=False,
+                        low_cpu_mem_usage=True
+                    )
+                    
+                    # Use CPU for inference to save memory
+                    self.model = self.model.to("cpu")
+                    
+                    # Enable memory efficient attention if available
+                    if hasattr(self.model, 'enable_attention_slicing'):
+                        self.model.enable_attention_slicing()
+                    
+                    # Enable memory efficient attention
+                    if hasattr(self.model, 'enable_vae_slicing'):
+                        self.model.enable_vae_slicing()
+                    
+                    # Enable xformers memory efficient attention if available
+                    try:
+                        if hasattr(self.model, 'enable_xformers_memory_efficient_attention'):
+                            self.model.enable_xformers_memory_efficient_attention()
+                    except:
+                        pass
+                except Exception as e:
+                    st.warning(f"Error loading image generation model: {str(e)}. Using fallback method.")
+                    self.model = None
                 
         return self.model
     
@@ -57,25 +63,140 @@ class ImageGenerator:
     def set_target_size(self, size):
         """Set the target image size"""
         self.target_size = size
+        
+    def set_aspect_ratio(self, aspect_ratio):
+        """Set the aspect ratio for generated images"""
+        self.aspect_ratio = aspect_ratio
+        
+        # Update target size based on aspect ratio while maintaining total pixels
+        base_pixels = self.target_size[0] * self.target_size[1]
+        
+        if aspect_ratio == "1:1":
+            # Square format
+            side = int(np.sqrt(base_pixels))
+            self.target_size = (side, side)
+        elif aspect_ratio == "16:9":
+            # Landscape format
+            width = int(np.sqrt(base_pixels * 16 / 9))
+            height = int(width * 9 / 16)
+            self.target_size = (width, height)
+        elif aspect_ratio == "9:16":
+            # Portrait format
+            height = int(np.sqrt(base_pixels * 16 / 9))
+            width = int(height * 9 / 16)
+            self.target_size = (width, height)
+    
+    def get_size_for_aspect_ratio(self, base_size, aspect_ratio):
+        """Calculate dimensions for a given aspect ratio while maintaining approximate total pixels"""
+        base_pixels = base_size[0] * base_size[1]
+        
+        if aspect_ratio == "1:1":
+            # Square format
+            side = int(np.sqrt(base_pixels))
+            return (side, side)
+        elif aspect_ratio == "16:9":
+            # Landscape format
+            width = int(np.sqrt(base_pixels * 16 / 9))
+            height = int(width * 9 / 16)
+            # Ensure dimensions are even numbers for video compatibility
+            width = width if width % 2 == 0 else width + 1
+            height = height if height % 2 == 0 else height + 1
+            return (width, height)
+        elif aspect_ratio == "9:16":
+            # Portrait format
+            height = int(np.sqrt(base_pixels * 16 / 9))
+            width = int(height * 9 / 16)
+            # Ensure dimensions are even numbers for video compatibility
+            width = width if width % 2 == 0 else width + 1
+            height = height if height % 2 == 0 else height + 1
+            return (width, height)
+        else:
+            # Default to original size
+            return base_size
     
     def generate_image(self, prompt, output_dir="temp"):
         """Generate a single image from a prompt"""
-        # Load the model if not already loaded
-        model = self.load_model()
-        
         # Ensure output directory exists
         os.makedirs(output_dir, exist_ok=True)
         
-        # Generate image with minimal inference steps to save resources
-        image = model(
-            prompt, 
-            num_inference_steps=self.inference_steps,
-            guidance_scale=7.5
-        ).images[0]
+        try:
+            # Load the model if not already loaded
+            model = self.load_model()
+            
+            if model is not None:
+                # Generate image with minimal inference steps to save resources
+                image = model(
+                    prompt, 
+                    num_inference_steps=self.inference_steps,
+                    guidance_scale=7.5
+                ).images[0]
+                
+                # Resize to target size for consistency and performance
+                if image.size != self.target_size:
+                    image = image.resize(self.target_size, Image.LANCZOS)
+            else:
+                # Fallback: Create a colored gradient image with text
+                from PIL import Image, ImageDraw, ImageFont, ImageFilter
+                
+                # Create a base image with gradient background
+                image = Image.new('RGB', self.target_size, color=(240, 240, 240))
+                draw = ImageDraw.Draw(image)
+                
+                # Create a gradient background
+                for y in range(image.height):
+                    for x in range(image.width):
+                        # Create a simple gradient
+                        r = int(200 + (x * 55 / image.width))
+                        g = int(200 + (y * 55 / image.height))
+                        b = 240
+                        draw.point((x, y), fill=(r, g, b))
+                
+                # Add some noise/texture
+                image = image.filter(ImageFilter.GaussianBlur(radius=1))
+                
+                # Add text from prompt (truncated)
+                draw = ImageDraw.Draw(image)
+                text = prompt[:50] + "..." if len(prompt) > 50 else prompt
+                
+                # Position text
+                text_width = draw.textlength(text, font=None)
+                text_position = ((image.width - text_width) / 2, image.height / 2)
+                
+                # Draw text
+                draw.text(text_position, text, fill=(0, 0, 0))
         
-        # Resize to target size for consistency and performance
-        if image.size != self.target_size:
-            image = image.resize(self.target_size, Image.LANCZOS)
+        except Exception as e:
+            st.warning(f"Error generating image: {str(e)}. Using fallback method.")
+            
+            # Fallback: Create a colored gradient image with text
+            from PIL import Image, ImageDraw, ImageFilter
+            
+            # Create a base image with gradient background
+            image = Image.new('RGB', self.target_size, color=(240, 240, 240))
+            draw = ImageDraw.Draw(image)
+            
+            # Create a gradient background
+            for y in range(image.height):
+                for x in range(image.width):
+                    # Create a simple gradient
+                    r = int(200 + (x * 55 / image.width))
+                    g = int(200 + (y * 55 / image.height))
+                    b = 240
+                    draw.point((x, y), fill=(r, g, b))
+            
+            # Add some noise/texture
+            image = image.filter(ImageFilter.GaussianBlur(radius=1))
+            
+            # Add text from prompt (truncated)
+            draw = ImageDraw.Draw(image)
+            text = prompt[:50] + "..." if len(prompt) > 50 else prompt
+            
+            # Position text
+            text_width = draw.textlength(text, font=None)
+            text_position = ((image.width - text_width) / 2, image.height / 2)
+            
+            # Draw text
+            draw.text(text_position, text, fill=(0, 0, 0))
         
         # Save the image
         image_path = f"{output_dir}/image_{int(time.time() * 1000)}.png"
@@ -85,9 +206,6 @@ class ImageGenerator:
     
     def generate_images(self, prompts, output_dir="temp", progress_callback=None, parallel=False, max_workers=4):
         """Generate images from the prompts"""
-        # Load the model if not already loaded
-        model = self.load_model()
-        
         # Ensure output directory exists
         os.makedirs(output_dir, exist_ok=True)
         
@@ -153,9 +271,6 @@ class ImageGenerator:
     
     def batch_generate_images(self, prompts, batch_size=2, output_dir="temp", progress_callback=None):
         """Generate images in batches to optimize memory usage"""
-        # Load the model if not already loaded
-        model = self.load_model()
-        
         # Ensure output directory exists
         os.makedirs(output_dir, exist_ok=True)
         
@@ -171,20 +286,7 @@ class ImageGenerator:
             # Generate images for this batch
             batch_images = []
             for j, prompt in enumerate(batch_prompts):
-                # Generate image
-                image = model(
-                    prompt, 
-                    num_inference_steps=self.inference_steps,
-                    guidance_scale=7.5
-                ).images[0]
-                
-                # Resize to target size
-                if image.size != self.target_size:
-                    image = image.resize(self.target_size, Image.LANCZOS)
-                
-                # Save the image
-                image_path = f"{output_dir}/image_{i+j}_{int(time.time() * 1000)}.png"
-                image.save(image_path)
+                image_path = self.generate_image(prompt, output_dir)
                 batch_images.append(image_path)
             
             # Add batch results to overall results