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

def generate_image(api_key, prompt, width, height, steps):
    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=width,
            height=height,
            steps=steps,
            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(scale=1):
            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
            )
            
            with gr.Row():
                width_input = gr.Number(
                    label="Width",
                    value=1024,
                    minimum=256,
                    maximum=2048,
                    step=64
                )
                height_input = gr.Number(
                    label="Height",
                    value=768,
                    minimum=256,
                    maximum=2048,
                    step=64
                )
                steps_input = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=50,
                    value=4,
                    step=1
                )
                
            generate_button = gr.Button("Generate Image")
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image")
    
    generate_button.click(
        fn=generate_image,
        inputs=[api_key_input, prompt_input, width_input, height_input, steps_input],
        outputs=image_output
    )
    
    gr.Markdown("""
    ## Instructions
    1. Enter your Together API Key (get one from https://www.together.ai)
    2. Enter a descriptive prompt in the text box
    3. Adjust the width, height, and steps as needed
        - Width and height control the dimensions of the generated image
        - Steps control the number of diffusion steps (higher = more detail but slower)
    4. Click "Generate Image" 
    5. 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()