import streamlit as st
import os
import requests
from PIL import Image
from io import BytesIO
import replicate
from stability_sdk import client
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation

# 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 resize_image(image_bytes, max_size=(256, 256)):
    # Open the image from bytes
    img = Image.open(BytesIO(image_bytes))
    
    # Resize the image
    img.thumbnail(max_size)
    
    # Save it back to bytes
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    return buffer.getvalue()

def upscale_image_stable_diffusion(image_bytes):
    # Set up environment variables
    os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
    os.environ['STABILITY_KEY'] = STABLE_DIFFUSION_API_KEY

    # Set up the connection to the API
    stability_api = client.StabilityInference(
        key=os.environ['STABILITY_KEY'],
        upscale_engine="stable-diffusion-x4-latent-upscaler",
        verbose=True,
    )

    # Open the image from bytes
    img = Image.open(BytesIO(image_bytes))

    # Call the upscale API
    answers = stability_api.upscale(init_image=img)

    # Process the response
    upscaled_img_bytes = None
    for resp in answers:
        for artifact in resp.artifacts:
            if artifact.type == generation.ARTIFACT_IMAGE:
                upscaled_img = Image.open(BytesIO(artifact.binary))
                upscaled_img_bytes = BytesIO()
                upscaled_img.save(upscaled_img_bytes, format='PNG')
                upscaled_img_bytes = upscaled_img_bytes.getvalue()
    
    return upscaled_img_bytes

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("Resizing image...")
        resized_image_bytes = resize_image(image_bytes)
        
        st.success("Upscaling image with stable-diffusion-x4-latent-upscaler...")
        # use the resized_image_bytes instead of the original image_bytes
        upscaled_image_bytes = upscale_image_stable_diffusion(resized_image_bytes)
        
        st.success("Further upscaling image with GFPGAN...")
        # Again use the upscaled_image_bytes from the previous step for the GFPGAN
        img = further_upscale_image(upscaled_image_bytes)
        
        st.image(img, caption='Upscaled Image', use_column_width=True)

if __name__ == "__main__":
    main()