from flask import Flask, jsonify, request
from pathlib import Path
import sys
import torch
import os
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler,  DiffusionPipeline
import streamlit as st

# model_path = WEIGHTS_DIR             # If you want to use previously trained model saved in gdrive, replace this with the full path of model in gdrive

pipe = StableDiffusionPipeline.from_pretrained("Divyanshu04/Finetuned-sd-vae", safety_checker=None, torch_dtype=torch.float32).to("cuda")
# pipe = DiffusionPipeline.from_pretrained("Divyanshu04/Finetuned-sd-vae", safety_checker=None, torch_dtype=torch.float32).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
g_cuda = None
         
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))

app = Flask(__name__)
  

# @app.route("/", methods=["POST"])
def generate():

    # prompt = request.form['prompt']
    # negative_prompt = request.form['Negative prompt']
    # num_samples = request.form['No. of samples']

    prompt = st.text_area(label = "prompt", key="pmpt")
    negative_prompt = st.text_area(label = "Negative prompt", key="ng_pmpt")
    num_samples = st.number_input("No. of samples")

    res = st.button("Generate", type="primary")

    if res:

        guidance_scale = 7.5
        num_inference_steps = 24
        height = 512
        width = 512

        g_cuda = torch.Generator(device='cuda')
        seed = 52362
        g_cuda.manual_seed(seed)

        # commandline_args = os.environ.get('COMMANDLINE_ARGS', "--skip-torch-cuda-test --no-half")

        with autocast("cuda"), torch.inference_mode():
            images = pipe(
                prompt,
                height=height,
                width=width,
                negative_prompt=negative_prompt,
                num_images_per_prompt=num_samples,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=g_cuda
            ).images
        
        return {"message": "successful"}
    
    else:
        st.write('<Enter parameters to generate image>')



  
# driver function
if __name__ == '__main__':
    generate()