import gradio as gr
import requests

# Replicate API token and model version
API_TOKEN = "r8_W4ccwSux3XNNLWzMEWJnRuAADU3x6oK2gWm45"  # Replace with your API token
MODEL_URL = "https://api.replicate.com/v1/predictions"
MODEL_VERSION = "613a21a57e8545532d2f4016a7c3cfa3c7c63fded03001c2e69183d557a929db"  # Replace with your model version ID

def generate_image(prompt):
    headers = {
        "Authorization": f"Token {API_TOKEN}",
        "Content-Type": "application/json"
    }
    
    data = {
        "version": MODEL_VERSION,
        "input": {
            "prompt": prompt
        }
    }
    
    response = requests.post(MODEL_URL, headers=headers, json=data)
    
    if response.status_code == 200:
        result = response.json()
        if "output" in result:
            image_url = result["output"][0]
            return image_url
        else:
            return "Error: No output in the response"
    else:
        return f"Error: {response.status_code} - {response.text}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Flux LoRA Image Generation")
    text_input = gr.Textbox(label="Enter a prompt", placeholder="e.g., A sunset over a mountain range")
    image_output = gr.Image(label="Generated Image")
    generate_button = gr.Button("Generate Image")
    generate_button.click(fn=generate_image, inputs=text_input, outputs=image_output)

if __name__ == "__main__":
    demo.launch()