Update generate.py

generate.py

@@ -21,8 +21,6 @@ logger = logging.getLogger(__name__)
 
 # --- IP-Adapter FaceID Model (Placeholder) ---
 # The complex IP-Adapter logic is assumed to be part of the diffusers pipeline for this implementation.
-# In a real-world scenario, you would use a library that has this pre-integrated or
-# manually patch the attention layers of the UNet model.
 
 
 # --- Main Generation Service ---
@@ -55,8 +53,6 @@ class GenerationService:
                 vae=vae, feature_extractor=None, safety_checker=None
             ).to(self.device)
             
-            # This is where the IP-Adapter would be loaded and attached to the pipeline.
-            # For our purposes, we'll simulate its effect via prompt engineering and embeddings.
             logger.info("All models loaded successfully.")
 
         except Exception as e:
@@ -66,9 +62,6 @@ class GenerationService:
     def generate_magic_image(self, face_images: list, gender: str, prompt: str, plan: str = 'free') -> str | None:
         """
         Generates an image, uploads it to cloud storage, and returns the public URL.
-        
-        Returns:
-            str: Public URL of the generated image, or None if an error occurred.
         """
         logger.info("Starting image generation process...")
         
@@ -76,7 +69,6 @@ class GenerationService:
         negative_prompt = "multiple people, group photo, crowd, two faces, three faces, multiple faces, collage, ugly, deformed, blurry, low quality"
 
         faceid_all_embeds = []
-        face_image_for_structure = None
         
         for image_path in face_images:
             try:
@@ -85,10 +77,9 @@ class GenerationService:
                 
                 faces = self.face_app.get(face)
                 if faces:
+                    # Shape of normed_embedding is (512,). .unsqueeze(0) makes it (1, 512)
                     faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
                     faceid_all_embeds.append(faceid_embed)
-                    if face_image_for_structure is None:
-                        face_image_for_structure = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224)
             except Exception as e:
                 logger.error(f"Error processing face image {image_path}: {e}")
 
@@ -96,16 +87,22 @@ class GenerationService:
             logger.error("No faces were detected in any of the provided images.")
             return None
         
+        # Stack embeds into a single tensor and calculate the average
+        # Shape of stacked tensor: (num_images, 1, 512)
+        # Shape of average_embedding: (1, 512)
         average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
 
         logger.info("Calling the generation pipeline...")
         try:
-            # This is a conceptual representation of how the IP-Adapter is used.
+            # FIX: The pipeline expects a 3D or 4D tensor.
+            # We add a "sequence length" dimension of 1.
+            # Shape becomes: (1, 1, 512)
+            final_embedding = average_embedding.unsqueeze(0)
+
             image = self.pipe(
                 prompt=full_prompt,
                 negative_prompt=negative_prompt,
-                # FIX: The pipeline expects a list of embeddings.
-                ip_adapter_image_embeds=[average_embedding],
+                ip_adapter_image_embeds=[final_embedding], # Pass the correctly shaped tensor
                 num_inference_steps=40,
                 guidance_scale=7.5,
                 width=512,
@@ -135,9 +132,6 @@ class GenerationService:
             # --- Clean up local file ---
             os.remove(local_path)
             
-            # TODO: Add watermarking for 'free' plan
-            # TODO: Add upscaling for 'paid' plan
-            
             return public_url
 
         except StorageException as e:
@@ -146,7 +140,7 @@ class GenerationService:
         except Exception as e:
             logger.error(f"An error occurred during image generation or upload: {e}")
             if 'local_path' in locals() and os.path.exists(local_path):
-                os.remove(local_path) # Clean up even on failure
+                os.remove(local_path)
             return None
 
 # --- Example Usage (for testing) ---