Refactor: Direct SDK call for image generation and cleanup

src/api_clients.py

@@ -27,8 +27,6 @@ from file_downloader import FileDownloader
 
 
 # --- NEW IMPORTS ---
-from google.oauth2 import service_account
-import vertexai
 from google_src.gcs_utils import get_gcs_client, get_gcs_credentials, upload_file_to_gcs, list_gcs_files, create_bucket_if_not_exists
 from google_src.setup_gcs_permissions import setup_bucket_permissions
 
@@ -101,75 +99,6 @@ class APIClients:
             await self.store_in_cache(file_path, f"{method_type}_{duration}", ".txt")
         except: pass
 
-    async def generate_image(self, prompt: str) -> Optional[str]:
-        """
-        Generate image using Vertex AI Imagen 4 Ultra
-
-        Args:
-            prompt: Image generation prompt
-
-        Returns:
-            Local path to generated image or None
-        """
-        try:
-            if os.getenv("TEST_AUTOMATION", "").lower() == "true":
-                return f"{os.getenv('TEST_DATA_DIRECTORY')}/image.png"
-
-            url = await self.get_cache_url("generate_image", ".jpg")
-            if url:
-                return url
-
-            logger.info(f"🎨 Generating image with Imagen 4 Ultra: {prompt[:200]}...")
-            image_path = ai_studio_sdk.generate_image(prompt)
-            if image_path:
-                await self.store_in_cache(image_path, "generate_image", ".jpg")
-                return image_path
-            raise ValueError("Image not Generated.")
-            from vertexai.preview.vision_models import ImageGenerationModel
-
-            vertex_ai_credentials = None
-            vertex_ai_project_id = self.config.get("gcp_project_id")  # Fallback to main config
-
-            # Check for the specific Vertex AI credentials in the environment variable
-            vertex_creds_json_str = os.getenv("VERTEX_AI_CREDENTIALS_JSON")
-            if vertex_creds_json_str:
-                logger.info("Found specific VERTEX_AI_CREDENTIALS_JSON. Initializing Vertex AI with these credentials.")
-                try:
-                    # Parse the JSON string from the environment variable
-                    creds_info = json.loads(vertex_creds_json_str)
-                    vertex_ai_credentials = service_account.Credentials.from_service_account_info(creds_info)
-                    # Get the project ID directly from the credentials, which is more reliable
-                    vertex_ai_project_id = creds_info.get("project_id")
-                except Exception as e:
-                    logger.error(f"❌ Failed to parse VERTEX_AI_CREDENTIALS_JSON: {e}. Falling back to default credentials.")
-            else:
-                logger.info("VERTEX_AI_CREDENTIALS_JSON not found. Initializing Vertex AI with default (WIF) credentials.")
-
-            # Initialize Vertex AI with either the specific credentials or default (None)
-            vertexai.init(project=vertex_ai_project_id, location="us-central1", credentials=vertex_ai_credentials)
-
-            model = ImageGenerationModel.from_pretrained("imagen-4.0-ultra-generate-001")
-            images = model.generate_images(
-                prompt=prompt,
-                number_of_images=1,
-                aspect_ratio="9:16",
-                safety_filter_level="block_some",
-                person_generation="allow_adult",
-            )
-            if not images:
-                logger.error("❌ Imagen API returned no images.")
-                return None
-
-            # Save to temp file
-            import tempfile
-            output_path = f"/tmp/hook_image_{hash(prompt)}.png"
-            images[0].save(location=output_path, include_generation_parameters=False)
-            logger.info(f"✓ Image generated with Imagen 4 Ultra (9:16): {output_path}")
-            return output_path
-
-        except Exception as e:
-            logger.error(f"❌ Imagen 4 Ultra generation failed: {e}")
-            raise
 
     async def generate_video(self, prompt: str, duration: int, image_input: str = None) -> Dict:
         """

src/automation.py

@@ -469,7 +469,7 @@ class ContentAutomation:
                 logger.info(f"Using provided captions: {captions[:50]}...")
 
             # Step 1: Generate image using Imagen 4 Ultra
-            image_path = await self.api_clients.generate_image(strategy["gemini_prompt"])
+            image_path = ai_studio_sdk.generate_image(strategy["gemini_prompt"])
             if not image_path:
                 raise Exception("Image generation failed")