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
import io
from PIL import Image
from huggingface_hub import login

# HF_TOKEN = os.environ.get("HF_TOKEN")

login(token='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')


pipe = StableDiffusionPipeline.from_pretrained("Divyanshu04/Finetuned-model", safety_checker=None, torch_dtype=torch.float16).to("cpu")

pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
# pipe.enable_xformers_memory_efficient_attention()  #if gpu is available
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():

    with st.form(key="Form :", clear_on_submit = True):
        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", step=1)

        Submit = st.form_submit_button(label='Submit')

    if Submit:

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

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


        with autocast("cpu"), 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
        
        st.image(images)
    
    else:
        st.write('<Enter parameters to generate image>')



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