from diffusers import StableDiffusionPipeline
import torch
import requests
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionImg2ImgPipeline

device = "cpu"
#model_path = "weights"
#model_id_or_path = "runwayml/stable-diffusion-v1-5"

model_id = "pwc-india/tartan_weights"
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

pipe2 = StableDiffusionImg2ImgPipeline.from_pretrained(model_id).to(device)

import gradio as gr

def generate_txt2img(prompt):
    return pipe(prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
def generate_img2img(img, prompt):
    image = Image.fromarray(img)
    return pipe2(prompt=prompt, image=image, strength=0.75, guidance_scale=7.5).images[0]

with gr.Blocks() as demo:
    with gr.Tab("Text2Image"):
        inp_txt = gr.Text(showlabel=False, placeholder="Enter your prompt here...")
        btn = gr.Button("Generate")
        out_img = gr.Image()
        btn.click(fn=generate_txt2img, inputs=[inp_txt], outputs=[out_img])
    with gr.Tab("Image2Image"):
        inp_img = gr.Image()
        inp_txt2 = gr.Text(showlabel=False,placeholder="Enter your prompt here...")
        btn2 = gr.Button("Generate")
        out_img2 = gr.Image()
        btn2.click(fn=generate_img2img, inputs=[inp_img, inp_txt2], outputs=[out_img2])
        
demo.launch(debug=True)