app.py

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
import threading

# 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):
    # Implement art generation logic using the Stable Diffusion API
    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",  # Set the Accept header to receive an image directly
    }

    json_data = {
        "inputs": prompt
    }

    while True:
        while True:
            response = requests.post('https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5', headers=headers, json=json_data)

            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))
            
            # Check if the image is black
            if not is_black_image(image):
                break

    return image


import threading

class ArtGenerationTimeout(Exception):
    pass

def generate_art_with_timeout(mood, timeout=120):
    result = {}

    def generate_and_store_result():
        result['art'] = generate_art(mood)

    thread = threading.Thread(target=generate_and_store_result)
    thread.start()
    thread.join(timeout)

    if thread.is_alive():
        raise ArtGenerationTimeout()
    else:
        return result['art']

def mood_art_generator(image):
    mood = get_mood_from_image(image)
    print("Mood:", mood)
    if mood:
        try:
            art = generate_art_with_timeout(mood)
            output_text = f"You seem to be {mood}. Here's an artwork representing it!"
            return art, output_text
        except ArtGenerationTimeout:
            return None, "Art generation took too long. Please try again."
    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()