import gradio as gr
import base64
import io
from PIL import Image
from together import Together

# Initialize the Together AI client
client = Together()

def generate_image(prompt):
    try:
        # Call the Together API to generate an image
        response = client.images.generate(
            prompt=prompt,
            model="black-forest-labs/FLUX.1-schnell-Free",
            width=1024,
            height=768,
            steps=4,
            n=1,
            response_format="b64_json"
        )
        
        # Get the base64 encoded image
        image_b64 = response.data[0].b64_json
        
        # Convert base64 to image
        image_data = base64.b64decode(image_b64)
        image = Image.open(io.BytesIO(image_data))
        
        return image
    except Exception as e:
        return f"Error generating image: {str(e)}"

# Create the Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# Together AI Image Generator")
    gr.Markdown("Generate images using the FLUX.1-schnell-Free model from Together AI")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Enter your prompt",
                placeholder="A beautiful sunset over mountains...",
                lines=3
            )
            generate_button = gr.Button("Generate Image")
        
        with gr.Column():
            image_output = gr.Image(label="Generated Image")
    
    generate_button.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=image_output
    )
    
    gr.Markdown("""
    ## Instructions
    1. Enter a descriptive prompt in the text box
    2. Click "Generate Image" 
    3. Wait for the image to be generated
    
    Note: You'll need to have your Together AI API key set up properly, 
    usually through environment variables or a config file.
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()