import requests
from PIL import Image
from io import BytesIO
import base64
import gradio as gr
from transformers import CLIPProcessor, CLIPModel
import numpy as np
import time

# Replace with your own API key
STABLE_DIFFUSION_API_KEY = "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"

# Load the CLIP model and processor
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

def get_mood_from_image(image: Image.Image):
    moods = ["scared", "angry", "happy", "sad", "disgusted", "surprised"]
    
    # Create unique prompts for each mood
    prompts = [
        "The emotion conveyed by this image is fear. The person looks scared and tense.",
        "The emotion conveyed by this image is anger. The person looks furious and irritated.",
        "The emotion conveyed by this image is happy. The person looks happy and cheerful.",
        "The emotion conveyed by this image is sadness. The person looks unhappy and gloomy.",
        "The emotion conveyed by this image is disgust. The person looks repulsed and sickened.",
        "The emotion conveyed by this image is surprise. The person looks astonished and amazed.",
    ]
    
    # Prepare the inputs for the model
    inputs = processor(text=prompts, images=image, return_tensors="pt", padding=True)
    
    # Run the model
    logits = model(**inputs).logits_per_image
    probs = logits.softmax(dim=-1).tolist()

    # Calculate the scores for each mood
    mood_scores = {}
    for mood, score in zip(moods, probs[0]):
        mood_scores[mood] = score
    print("Mood Scores:", mood_scores)
    # Select the mood with the highest score
    selected_mood = max(mood_scores, key=mood_scores.get)

    return selected_mood
def is_black_image(image: Image.Image) -> bool:
    img_array = np.array(image)
    return np.all(img_array == 0)

def generate_art(mood, max_retries=3, request_timeout=30):
    prompt = f"{mood} generative art with vibrant colors and intricate patterns ({str(np.random.randint(1, 10000))})"
    
    headers = {
        "Authorization": f"Bearer {STABLE_DIFFUSION_API_KEY}",
        "Accept": "image/jpeg",
    }

    json_data = {
        "inputs": prompt
    }

    retries = 0
    while retries < max_retries:
        try:
            response = requests.post('https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5', headers=headers, json=json_data, timeout=request_timeout)
        except requests.exceptions.Timeout:
            print(f"Request timed out after {request_timeout} seconds. Retrying...")
            retries += 1
            continue

        if response.status_code == 503:
            print("Model is loading, waiting for 30 seconds before retrying...")
            time.sleep(30)
            continue

        if response.status_code != 200:
            print(f"Error: API response status code {response.status_code}")
            print("Response content:")
            print(response.content)
            return None

        image = Image.open(BytesIO(response.content))

        if not is_black_image(image):
            break

        retries += 1

    if retries == max_retries:
        return None

    return image

def mood_art_generator(image):
    mood = get_mood_from_image(image)
    print("Mood:", mood)
    if mood:
        art = generate_art(mood)
        output_text = f"You seem to be {mood}. Here's an artwork representing it!"
        return art, output_text
    else:
        return None, "Failed to generate artwork."

iface = gr.Interface(
    fn=mood_art_generator,
    inputs=gr.inputs.Image(shape=(224, 224), image_mode="RGB", source="upload"),
    outputs=[gr.outputs.Image(type="pil"), gr.outputs.Textbox()],
    title="Mood-based Art Generator",
    description="Upload an image of yourself and let the AI generate artwork based on your mood.",
    allow_flagging=False,
    analytics_enabled=False,
    share=True
)
iface.launch()