import gradio as gr
from PIL import Image
from diffusers.utils import load_image
import os, random, gc, re, json, time, shutil, glob
import PIL.Image
import tqdm
from accelerate import Accelerator
from huggingface_hub import HfApi, InferenceClient, ModelCard, RepoCard, upload_folder, hf_hub_download, HfFileSystem
HfApi=HfApi()
HF_TOKEN=os.getenv("HF_TOKEN")
HF_HUB_DISABLE_TELEMETRY=1
DO_NOT_TRACK=1
HF_HUB_ENABLE_HF_TRANSFER=0
accelerator = Accelerator(cpu=True)
InferenceClient=InferenceClient()

apol=[]

hf_hub_download(repo_id="black-forest-labs/FLUX.1-schnell", filename="ae.safetensors", local_dir=".")
vae_path = './ae.safetensors'
hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.safetensors", local_dir=".")
clip_path = './clip_l.safetensors'
hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir=".")
t5xxl_path = './t5xxl_fp16.safetensors'
hf_hub_download(repo_id="enhanceaiteam/MidFlow", filename="midflow_fp16_low.safetensors", local_dir=".")
model_path = './midflow_fp16_low.safetensors'

from stable_diffusion_cpp import StableDiffusion
flux_diffusion = StableDiffusion(diffusion_model_path=model_path,clip_l_path=clip_path,t5xxl_path=t5xxl_path,vae_path=vae_path,)
def chdr(apol,prompt,modil,stips,fnamo,gaul):
    try:
        type="sd_cpp"
        los=""
        tre='./tmpo/'+fnamo+'.json'
        tra='./tmpo/'+fnamo+'_0.png'
        ##trm='./tmpo/'+fnamo+'_1.png'
        flng=["yssup", "sllab", "stsaerb", "sinep", "selppin", "ssa", "tnuc", "mub", "kcoc", "kcid", "anigav", "dekan", "edun", "slatineg", "xes", "nrop", "stit", "ttub", "bojwolb", "noitartenep", "kcuf", "kcus", "kcil", "elttil", "gnuoy", "thgit", "lrig", "etitep", "dlihc", "yxes"]
        flng=[itm[::-1] for itm in flng]
        ptn = r"\b" + r"\b|\b".join(flng) + r"\b"
        if re.search(ptn, prompt, re.IGNORECASE):
            print("onon buddy")
        else:
            dobj={'img_name':fnamo,'model':modil,'lora':los,'prompt':prompt,'steps':stips,'type':type}
            with open(tre, 'w') as f:
                json.dump(dobj, f)
            HfApi.upload_folder(repo_id="JoPmt/hf_community_images",folder_path="./tmpo",repo_type="dataset",path_in_repo="./",token=HF_TOKEN)
        dobj={'img_name':fnamo,'model':modil,'lora':los,'prompt':prompt,'steps':stips,'type':type,'haed':gaul,}
        try:
            for pxn in glob.glob('./tmpo/*.png'):
                os.remove(pxn)
        except:
            print("lou")
        with open(tre, 'w') as f:
            json.dump(dobj, f)
        HfApi.upload_folder(repo_id="JoPmt/Tst_datast_imgs",folder_path="./tmpo",repo_type="dataset",path_in_repo="./",token=HF_TOKEN)
        try:
            for pgn in glob.glob('./tmpo/*.png'):
                os.remove(pgn)
            for jgn in glob.glob('./tmpo/*.json'):
                os.remove(jgn)
            del tre
            del tra
            ##del trm
        except:
            print("cant")
    except:
        print("failed to make obj")

def plax(gaul,req: gr.Request):
    gaul=str(req.headers)
    return gaul

def plex(prompt,neg_prompt,stips,nut,wit,het,gaul,progress=gr.Progress(track_tqdm=True)):
    gc.collect()
    apol=[]
    modil="flux/flux-schnell"
    fnamo=""+str(int(time.time()))+""
    if nut == 0:
        nm = random.randint(1, 2147483616)
        while nm % 32 != 0:
            nm = random.randint(1, 2147483616)
    else:
        nm=nut
    image = flux_diffusion.txt_to_img(prompt=str(prompt),negative_prompt=str(neg_prompt),sample_steps=int(stips),seed=int(nm),height=int(het),width=int(wit),cfg_scale=1.0,sample_method="euler",)
    for a, imze in enumerate(image):
        apol.append(imze)
        imze.save('./tmpo/'+fnamo+'_'+str(a)+'.png', 'PNG')
    chdr(apol,prompt,modil,stips,fnamo,gaul)
    return apol

def aip(ill,api_name="/run"):
    return
def pit(ill,api_name="/predict"):
    return

info="""
    <div>
    <text>
    This is a demo of Flux.1-schnell on CPU.
    This space uses stable-diffusion.cpp by leejet.
    </text>
    <img src=''></img>
    </div>
"""

with gr.Blocks(theme=random.choice([gr.themes.Monochrome(),gr.themes.Base.from_hub("gradio/seafoam"),gr.themes.Base.from_hub("freddyaboulton/dracula_revamped"),gr.themes.Glass(),gr.themes.Base(),]),analytics_enabled=False) as iface:
    ##iface.description="Running on cpu, very slow! by JoPmt."
    out=gr.Gallery(label="Generated Output Image", columns=1)
    inut=gr.Textbox(label="Prompt",value="")
    gaul=gr.Textbox(visible=False)
    gr.HTML(info)
    btn=gr.Button("GENERATE")
    with gr.Accordion("Advanced Settings", open=False):
        inet=gr.Textbox(label="Negative_prompt", value="lowres,text,bad quality,low quality,jpeg artifacts,ugly,bad hands,bad face,blurry,bad eyes,watermark,signature")
        inyt=gr.Slider(label="Num inference steps",minimum=1,step=1,maximum=10,value=3)
        indt=gr.Slider(label="Manual seed (leave 0 for random)",minimum=0,step=32,maximum=2147483616,value=0)
        inwt=gr.Slider(label="Width",minimum=256,step=32,maximum=1024,value=512)
        inht=gr.Slider(label="Height",minimum=256,step=32,maximum=1024,value=512)
    
    btn.click(fn=plax,inputs=gaul,outputs=gaul).then(fn=plex, outputs=[out], inputs=[inut,inet,inyt,indt,inwt,inht,gaul])

iface.queue(max_size=1,api_open=False)
iface.launch(max_threads=20,inline=False,show_api=False)