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

def generate_image(api_key, prompt):
    try:
        # Initialize the Together AI client with the provided API key
        client = Together(api_key=api_key)
        
        # 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():
            api_key_input = gr.Textbox(
                label="Together API Key",
                placeholder="Enter your Together API key here...",
                type="password"  # Mask the API key for security
            )
            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=[api_key_input, prompt_input],
        outputs=image_output
    )
    
    gr.Markdown("""
    ## Instructions
    1. Enter your Together AI API key (get one from https://www.together.ai)
    2. Enter a descriptive prompt in the text box
    3. Click "Generate Image" 
    4. Wait for the image to be generated
    
    Note: Your API key is not stored and is only used for the current session.
    """)

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