moderation added

.gitignore

@@ -1,2 +1,4 @@
 .env
-venv/
+venv/
+.gitattributes
+__pycache__/

app.py

@@ -13,6 +13,8 @@ import mimetypes
 from workflow_handler import WanVideoWorkflow
 from video_config import MODEL_FRAME_RATES, calculate_frames
 import asyncio
+from openai import OpenAI
+import base64
 
 dotenv.load_dotenv()
 
@@ -137,6 +139,9 @@ supabase: Client = create_client(
    
 )
 
+# Initialize OpenAI client
+openai_client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+
 def initialize_gcs():
     """Initialize Google Cloud Storage client with credentials from environment"""
     try:
@@ -353,6 +358,122 @@ def poll_generation_status(generation_id):
         print(f"Error polling generation status: {e}")
         raise e
 
+async def moderate_prompt(prompt: str) -> dict:
+    """
+    Check if a text prompt contains NSFW content
+    """
+    try:
+        response = openai_client.moderations.create(input=prompt)
+        result = response.results[0]
+        
+        if result.flagged:
+            # Find which categories were flagged
+            flagged_categories = [
+                category for category, flagged in result.categories.model_dump().items() 
+                if flagged
+            ]
+            
+            return {
+                "isNSFW": True,
+                "reason": f"Content flagged for: {', '.join(flagged_categories)}"
+            }
+        
+        return {"isNSFW": False, "reason": None}
+    except Exception as e:
+        print(f"Error during prompt moderation: {e}")
+        return {"isNSFW": False, "reason": None}
+
+async def moderate_image(image_path: str) -> dict:
+    """
+    Check if an image contains NSFW content using OpenAI's vision capabilities
+    """
+    try:
+        # Convert image to base64
+        with open(image_path, "rb") as image_file:
+            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
+        
+        response = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a content moderation system. Your task is to determine if an image contains NSFW content. Respond with only 'NSFW' if the image contains inappropriate sexual content such as nudity, pornography (especially child pornography), or other explicit material. Otherwise, respond with 'SAFE'."
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Is this image appropriate or does it contain NSFW content?"},
+                        {"type": "image_url", "image_url": {
+                            "url": f"data:image/jpeg;base64,{base64_image}"
+                        }}
+                    ]
+                }
+            ],
+            max_tokens=10
+        )
+        
+        result = response.choices[0].message.content.strip()
+        
+        if "NSFW" in result:
+            return {
+                "isNSFW": True,
+                "reason": "Image contains potentially inappropriate content"
+            }
+        
+        return {"isNSFW": False, "reason": None}
+    except Exception as e:
+        print(f"Error during image moderation: {e}")
+        return {"isNSFW": False, "reason": None}
+
+async def moderate_combined(prompt: str, image_path: str) -> dict:
+    """
+    Check if both a prompt and image contain NSFW content together
+    """
+    try:
+        # Convert image to base64
+        with open(image_path, "rb") as image_file:
+            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
+        
+        response = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a content moderation system. Your task is to determine if an image contains NSFW content. The image + text prompt combined will be generated into a video. Please assess if the generated video would be inappropriate in terms of sexual content only (pornography or child pornography). Respond with 'NSFW: reason' if inappropriate, or 'SAFE' if appropriate."
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f'Please moderate this image and prompt combination for an image-to-video generation:\n\nPrompt: "{prompt}"'
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{base64_image}"
+                            }
+                        }
+                    ]
+                }
+            ],
+            max_tokens=150
+        )
+        
+        result = response.choices[0].message.content.strip()
+        if result.startswith("NSFW:"):
+            return {
+                "isNSFW": True,
+                "reason": result[5:].strip()
+            }
+        return {
+            "isNSFW": False,
+            "reason": None
+        }
+    except Exception as e:
+        print(f"Error during combined moderation: {e}")
+        return {"isNSFW": False, "reason": None}
+
 async def generate_video(input_image, subject, duration, selected_index, progress=gr.Progress()):
     try:
         # Initialize workflow handler with explicit paths
@@ -451,6 +572,24 @@ async def handle_generation(image_input, subject, duration, selected_index, prog
     try:
         if selected_index is None:
             raise gr.Error("You must select a LoRA before proceeding.")
+        
+        # First, moderate the prompt
+        prompt_moderation = await moderate_prompt(subject)
+        print(f"Prompt moderation result: {prompt_moderation}")
+        if prompt_moderation["isNSFW"]:
+            raise gr.Error(f"Content moderation failed: {prompt_moderation['reason']}")
+        
+        # Then, moderate the image
+        image_moderation = await moderate_image(image_input)
+        print(f"Image moderation result: {image_moderation}")
+        if image_moderation["isNSFW"]:
+            raise gr.Error(f"Content moderation failed: {image_moderation['reason']}")
+        
+        # Finally, check the combination
+        combined_moderation = await moderate_combined(subject, image_input)
+        print(f"Combined moderation result: {combined_moderation}")
+        if combined_moderation["isNSFW"]:
+            raise gr.Error(f"Content moderation failed: {combined_moderation['reason']}")
             
         # Generate the video and get generation ID
         generation_id = await generate_video(image_input, subject, duration, selected_index)