#___________________________________________________________________________________________________________________________

import streamlit as st 
import os

#___________________________________________________________________________________________________________________________

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
from datasets import load_dataset
from PIL import Image  
import re

#___________________________________________________________________________________________________________________________

st.title('IMGTEXTA')
prompt=st.text_input('Enter Your Prompt')
Generate=st.button('Generate')

#___________________________________________________________________________________________________________________________

model_id = "CompVis/stable-diffusion-v1-4"
device = "cpu"
st.info('1')
#___________________________________________________________________________________________________________________________

pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token='Add Your Info', torch_dtype=torch.float32)
def dummy(images, **kwargs): return images, False 
pipe.safety_checker = dummy
st.info('2')
#___________________________________________________________________________________________________________________________

def infer(prompt, width, height, steps, scale, seed):      
    if seed == -1:	
        images_list = pipe(
            [prompt],
            height=height, 
            width=width,
            num_inference_steps=steps,
            guidance_scale=scale,
            generator=torch.Generator(device=device).manual_seed(seed))
    else:
        images_list = pipe(
            [prompt],
            height=height, 
            width=width,
            num_inference_steps=steps,
            guidance_scale=scale)

    return images_list["sample"]
    st.info('5')
st.info('3')
#___________________________________________________________________________________________________________________________

def onclick(prompt):
    st.image(infer(prompt,512,512,30,7.5,-1))



if Generate==True:
    onclick(prompt)
    st.info('4')

#___________________________________________________________________________________________________________________________