app.py

import os
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoProcessor, FocalNetForImageClassification
import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
from huggingface_hub import InferenceClient
import requests
from io import BytesIO

# Paths and model setup
image_folder = "path_to_your_image_folder"  # Specify the path to your image folder
model_path = "MichalMlodawski/nsfw-image-detection-large"

# List of jpg files in the folder
jpg_files = [file for file in os.listdir(image_folder) if file.lower().endswith(".jpg")]

if not jpg_files:
    print("🚫 No jpg files found in folder:", image_folder)
    exit()

# Load the model and feature extractor
feature_extractor = AutoProcessor.from_pretrained(model_path)
model = FocalNetForImageClassification.from_pretrained(model_path)
model.eval()

# Image transformations
transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Mapping from model labels to NSFW categories
label_to_category = {
    "LABEL_0": "Safe",
    "LABEL_1": "Questionable",
    "LABEL_2": "Unsafe"
}

# Device configuration
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the diffusion pipeline
if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
    pipe.enable_xformers_memory_efficient_attention()
    pipe = pipe.to(device)
else: 
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
    pipe = pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

# Initialize the InferenceClient
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Inference function for generating images
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt=prompt, 
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale, 
        num_inference_steps=num_inference_steps, 
        width=width, 
        height=height,
        generator=generator
    ).images[0] 
    
    return image

# Respond function for the chatbot
def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}]
    
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    
    response = client.chat_completion(
        messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    
    return response.choices[0].message['content']

# Function to generate posts
def generate_post(prompt, max_tokens, temperature, top_p):
    response = client.chat_completion(
        [{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    return response.choices[0].message['content']

# Function to moderate posts
def moderate_post(post):
    # Implement your post moderation logic here
    if "inappropriate" in post:
        return "Post does not adhere to community guidelines."
    return "Post adheres to community guidelines."

# Function to generate images using the diffusion pipeline
def generate_image(prompt):
    generator = torch.manual_seed(random.randint(0, MAX_SEED))
    image = pipe(prompt=prompt, generator=generator).images[0]
    return image

# Function to moderate images
def moderate_image(image):
    # Convert the PIL image to a format that can be sent for moderation
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    image_bytes = buffered.getvalue()
    
    # Replace with your actual image moderation API endpoint
    moderation_api_url = "https://example.com/moderation/api"
    
    # Send the image to the moderation API
    response = requests.post(moderation_api_url, files={"file": image_bytes})
    result = response.json()
    
    # Check the result from the moderation API
    if result.get("moderation_status") == "approved":
        return "Image adheres to community guidelines."
    else:
        return "Image does not adhere to community guidelines."

# Create the Gradio interface
css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

with gr.Blocks(css=css) as demo:
    gr.Markdown("# AI-driven Content Generation and Moderation Bot")
    gr.Markdown(f"Currently running on {power_device}.")

    with gr.Tabs():
        with gr.TabItem("Chat"):
            with gr.Column():
                chat_interface = gr.ChatInterface(
                    respond,
                    additional_inputs=[
                        gr.Textbox(value="You are a friendly Chatbot meant to assist users in managing social media posts ensuring they meet community guidelines", label="System message", visible=False),
                        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens", visible=False),
                        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature", visible=False),
                        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)", visible=False),
                    ],
                )
                advanced_button = gr.Button("Show Advanced Settings")
                advanced_settings = gr.Column(visible=False)
                with advanced_settings:
                    chat_interface.additional_inputs[0].visible = True
                    chat_interface.additional_inputs[1].visible = True
                    chat_interface.additional_inputs[2].visible = True
                    chat_interface.additional_inputs[3].visible = True
                
                def toggle_advanced_settings():
                    advanced_settings.visible = not advanced_settings.visible
                
                advanced_button.click(toggle_advanced_settings, [], advanced_settings)
        
        with gr.TabItem("Generate Post"):
            post_prompt = gr.Textbox(label="Post Prompt")
            max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
            temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
            generate_button = gr.Button("Generate Post")
            generated_post = gr.Textbox(label="Generated Post")
            generate_button.click(generate_post, [post_prompt, max_tokens, temperature, top_p], generated_post)
        
        with gr.TabItem("Moderate Post"):
            post_content = gr.Textbox(label="Post Content")
            moderate_button = gr.Button("Moderate Post")
            moderation_result = gr.Textbox(label="Moderation Result")
            moderate_button.click(moderate_post, post_content, moderation_result)
        
        with gr.TabItem("Generate Image"):
            image_prompt = gr.Textbox(label="Image Prompt")
            generate_image_button = gr.Button("Generate Image")
            generated_image = gr.Image(label="Generated Image")
            generate_image_button.click(generate_image, image_prompt, generated_image)
        
        with gr.TabItem("Moderate Image"):
            uploaded_image = gr.Image(label="Upload Image")
            moderate_image_button = gr.Button("Moderate Image")
            image_moderation_result = gr.Textbox(label="Image Moderation Result")
            moderate_image_button.click(moderate_image, uploaded_image, image_moderation_result)

        with gr.TabItem("NSFW Classification"):
            selected_image = gr.Image(type="pil", label="Upload Image for NSFW Classification")
            classify_button = gr.Button("Classify Image")
            classification_result = gr.Textbox(label="Classification Result")

            def classify_nsfw(image):
                image_tensor = transform(image).unsqueeze(0)
                inputs = feature_extractor(images=image, return_tensors="pt")
                with torch.no_grad():
                    outputs = model(**inputs)
                    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
                    confidence, predicted = torch.max(probabilities, 1)
                label = model.config.id2label[predicted.item()]
                category = label_to_category.get(label, "Unknown")
                return f"Label: {label}, Category: {category}, Confidence: {confidence.item() * 100:.2f}%"

            classify_button.click(classify_nsfw, selected_image, classification_result)

demo.launch()