# main.py

import streamlit as st
import os
import requests
from PIL import Image
from io import BytesIO
import replicate

# Configure your API keys here
CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'

# Set up environment variable for Replicate API Token
os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token

def generate_image_from_text(prompt):
    r = requests.post('https://clipdrop-api.co/text-to-image/v1',
        files = {
            'prompt': (None, prompt, 'text/plain')
        },
        headers = { 'x-api-key': CLIPDROP_API_KEY }
    )
    
    if r.ok:
        return r.content
    else:
        r.raise_for_status()

def upscale_image_stable_diffusion(image_bytes):
    url = 'https://stable-diffusion-x4-latent-upscaler.com/v1/upscaling' # Update this with correct API endpoint
    headers = { 'x-api-key': STABLE_DIFFUSION_API_KEY }
    files = {
        'image': ('image.png', image_bytes, 'image/png')
    }
    
    r = requests.post(url, headers=headers, files=files)
    
    if r.ok:
        return r.content
    else:
        r.raise_for_status()

def further_upscale_image(image_bytes):
    # Run the GFPGAN model
    output = replicate.run(
        "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
        input={"img": BytesIO(image_bytes), "version": "v1.4", "scale": 16}
    )
    
    # The output is a URI of the processed image
    # We will retrieve the image data and save it
    response = requests.get(output)
    img = Image.open(BytesIO(response.content))
    img.save("upscaled.png")  # Save the upscaled image
    return img

def main():
    st.title("Image Generation and Upscaling")
    st.write("Enter a text prompt and an image will be generated and upscaled.")

    prompt = st.text_input("Enter a textual prompt to generate an image...")
    
    if prompt:
        st.success("Generating image from text prompt...")
        image_bytes = generate_image_from_text(prompt)
        
        st.success("Upscaling image with stable-diffusion-x4-latent-upscaler...")
        upscaled_image_bytes = upscale_image_stable_diffusion(image_bytes)
        
        st.success("Further upscaling image with GFPGAN...")
        img = further_upscale_image(upscaled_image_bytes)
        
        st.image(img, caption='Upscaled Image', use_column_width=True)

if __name__ == "__main__":
    main()