import os
from gradio import Interface
from PIL import Image
from groq import Groq
from diffusers import StableDiffusionPipeline
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Groq API setup
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Stable Diffusion setup using Hugging Face's diffusers library
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

# Function to generate an image and interact with Groq's API
def generate_and_interact(prompt):
    # Step 1: Generate image using Stable Diffusion
    image = model(prompt).images[0]
    
    # Step 2: Save the generated image (optional, for further use)
    image_path = "generated_image.png"
    image.save(image_path)

    # Step 3: Interact with Groq API (example: send metadata or a description of the image)
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": f"Analyze the generated image from this prompt: {prompt}"}],
        model="llama3-8b-8192",
        stream=False,
    )
    
    # Return the Groq response and the generated image
    return chat_completion.choices[0].message.content, image

# Gradio interface setup
iface = Interface(fn=generate_and_interact, inputs="text", outputs=["text", "image"])

# Launch the app
iface.launch()