#import streamlit as st
import gradio as gr
import torch
from PIL import Image
import numpy as np
from io import BytesIO
from diffusers import StableDiffusionImg2ImgPipeline
import os
import random

device="cpu"

USER_TOKEN=os.environ.get('HF_TOKEN_SD')

#1.4
#pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token = USER_TOKEN)

#1.5
#model_id = "runwayml/stable-diffusion-v1-5"
#pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token = USER_TOKEN)
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token = USER_TOKEN)

pipe.to(device)

source_img = gr.Image(source="upload", type="filepath", label="Landscape image 16:9")
#source_img = gr.Webcam(type="filepath")

def resize(rwidth,rheight,img):
  img = Image.open(img)
  img = img.resize((rwidth,rheight))
  return img

def infer(source_img):

    #Permanent settings here
    
    prompt = "abandoned buildings, empty streetscapes, surrounded by lush green vegetation, ground-level view, rusted steel, derelict, urban exploration, ruin, deserted, broken windows, burnt out, graffitti, ramshackle, homeless shelter, decay, 8k, photorealistic, hyper detailed"
    guide = 10
    steps = 40
    diffstr = 0.3
    seed = random.randint(0,2147483647)

    generator = torch.Generator("cpu").manual_seed(seed)     

    source_image = resize(1024,576, source_img)
    source_image.save('source.png')
        
    #1.4 image = pipe([prompt] * 1, init_image=source_image, strength=diffstr, guidance_scale=guide, num_inference_steps=steps).images[0]
    #1.5
    image = pipe([prompt] * 1, image=source_image, strength=diffstr, guidance_scale=guide, num_inference_steps=steps).images[0]
    
    #images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps)
    #images = []
    #images.append(image)
    
    return image

#gr.Interface(fn=infer, inputs=gr.Image(source="upload", type="filepath", label="input"), outputs=image, title = "Ruins of Tomorrow", description = "Upload an image of a place in a city in landscape format. Processing may take 5-15 minutes, please be patient!", article = "").queue(max_size=10).launch(enable_queue=True, debug=True)

title = "Ruins of Tomorrow"
description = "Upload an image of a place in a city in landscape format. Processing may take 5-15 minutes, please be patient!"

#ruins = gr.Interface(
#    fn=infer, 
#    inputs=[source_img],
#    outputs=gallery)

ruins = gr.Interface(
    title=title,
    description=description,
    fn=infer, 
    inputs=source_img,
    outputs="image")

ruins.queue(max_size=10)
ruins.launch(enable_queue=True)