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

device="cpu"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token = st.secrets['USER_TOKEN'])
pipe.to(device)

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

def infer(source_img, prompt, guide, steps, seed, Strength): 
    generator = torch.Generator("cpu").manual_seed(seed)     
    source_image = resize(512, source_img)
    source_image.save('source.png')
    image = pipe([prompt], init_image=source_image, strength=Strength, guidance_scale=guide, num_inference_steps=steps).images[0]
    return image

gr.Interface(fn=infer, inputs=[gr.Image(source="upload", type="filepath", label="Raw Image"), gr.Textbox(label = 'Prompt Input Text'),
    gr.Slider(2, 15, value = 7, label = 'Guidence Scale'),
    gr.Slider(10, 50, value = 25, step = 1, label = 'Number of Iterations'),
    gr.Slider(
        label = "Seed",
        minimum = 0,
        maximum = 2147483647,
        step = 1,
        randomize = True), gr.Slider(label='Strength', minimum = 0, maximum = 1, step = .05, value = .5)
    ], outputs='image', title = "Stable Diffusion Image to Image Pipeline CPU", description = "Upload an Image (must be .PNG and 512x512-2048x2048) enter a Prompt, or let it just do its Thing, then click submit. 10 Iterations takes about 300 seconds currently. For more informationon about Stable Diffusion or Suggestions for prompts, keywords, artists or styles see https://github.com/Maks-s/sd-akashic", article = "Code Monkey: <a href=\"https://huggingface.co/Manjushri\">Manjushri</a>").queue(max_size=10).launch(enable_queue=True, debug=True)