Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from random import randint
 from operator import itemgetter
 import bisect
-from all_models import tags_plus_models,models,models_plus_tags
+from all_models import tags_plus_models,models,models_plus_tags,find_warm_model_list
 from datetime import datetime
 from externalmod import gr_Interface_load
 import asyncio
@@ -13,7 +13,7 @@ lock = RLock()
 HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
 
 
-nb_req_simult=75 ########
+nb_req_simult=10 ########
 nb_gallery_model=5
 
 
@@ -21,7 +21,8 @@ tempo_update_actu=3.0
 #incr_update_actu={}
 
 now2 = 0
-inference_timeout = 200
+inference_timeout = 300
+inference_timeout_w = 70
 MAX_SEED = 2**32-1
 
 
@@ -34,6 +35,7 @@ cache_id_image={}
 cache_list_task={}
 cache_text_actu={}
 from_reload={}
+cache_list_task_w={}
 
 
 def load_fn(models):
@@ -117,13 +119,13 @@ async def infer(model_str, prompt, nprompt="", height=None, width=None, steps=No
             image = str(Path(png_path).resolve())
         return image
     return None
-def gen_fn(model_str, prompt, nprompt="", height=None, width=None, steps=None, cfg=None, seed=-1):
+def gen_fn(model_str, prompt, nprompt="", height=None, width=None, steps=None, cfg=None, seed=-1,timeout=inference_timeout):
     if model_str == 'NA':
         return None
     try:
         loop = asyncio.new_event_loop()
         result = loop.run_until_complete(infer(model_str, prompt, nprompt,
-                                         height, width, steps, cfg, seed, inference_timeout))
+                                         height, width, steps, cfg, seed, timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
         print(f"Task aborted: {model_str}")
@@ -166,7 +168,7 @@ def reset_cache_image_all_sessions(cache_image=cache_image):
         listT.clear()
     return 
     
-def set_session(id):
+def set_session(id,warm=False):
     if id==0:
         randTemp=randint(1,MAX_SEED)
         cache_image[f"{randTemp}"]=[]
@@ -174,7 +176,12 @@ def set_session(id):
         
         cache_list_task[f"{randTemp}"]=[]
         cache_text_actu[f"{randTemp}"]={}
+        if warm:
+            cache_text_actu[f"{randTemp}"]["warm"]=True
+        else:
+            cache_text_actu[f"{randTemp}"]["warm"]=False
         from_reload[f"{randTemp}"]=False
+        cache_list_task_w[f"{randTemp}"]=[]
         #incr_update_actu[f"{randTemp}"]=0
         return gr.Number(visible=False,value=randTemp)
     else :
@@ -307,6 +314,28 @@ def fonc_load_info(nb_of_models_to_gen,index_tag,index_first_model):
         list_models_temp.append(tags_plus_models[index_tag][2][i+index_first_model])
         str_temp+=f"\"{tags_plus_models[index_tag][2][i+index_first_model]}\",\n"
     return nb_of_models_to_gen,gr.Textbox(str_temp),gr.Dropdown(choices=[["",list_models_temp]], value=list_models_temp )
+
+def load_random_models(nb_of_models_to_gen,index_tag):
+    str_temp=""
+    list_models_temp=[]
+    list_random=[]
+    if nb_of_models_to_gen>=len(tags_plus_models[index_tag][2]):
+        str_temp+="warning : to many model chosen"
+        nb_of_models_to_gen= len(tags_plus_models[index_tag][2])
+        str_temp+=f" - only {nb_of_models_to_gen} will be use\n\n"
+        list_models_temp=tags_plus_models[index_tag][2]
+        for m in list_models_temp:
+            str_temp+=f"\"{m}\",\n"
+    else :
+        list_random=tags_plus_models[index_tag][2].copy()
+        for i in range(nb_of_models_to_gen):
+            i_rand=randint(0,len(list_random)-1)
+            m=list_random.pop(i_rand)
+            list_models_temp.append(m)
+            str_temp+=f"\"{m}\",\n"
+    return nb_of_models_to_gen,gr.Textbox(str_temp),gr.Dropdown(choices=[["",list_models_temp]], value=list_models_temp )      
+        
+
 def fonc_load_info_custom(nb_of_models_to_gen,list_model_custom,index_first_model):
     str_temp=""
     list_models_temp=[]
@@ -342,7 +371,7 @@ def crea_list_task(id_session,list_param,list_models_to_gen,nb_images_by_prompt)
                 dict_temp["task"].append({"prompt":p[0],"nprompt":p[1],"width":p[2],"height":p[3],"steps":p[4],"cfg":p[5],"seed":p[6]})
         cache_list_task[f"{id_session}"].append(dict_temp)
     
-    cache_text_actu[f"{id_session}"]={"nb_modules_use":nb_req_simult,"stop":False,"nb_fail":0,
+    cache_text_actu[f"{id_session}"]={"nb_modules_use":nb_req_simult,"stop":False,"nb_fail":0,"warm":False,
                                       "nb_models_to_do":len(list_models_to_gen) ,"nb_models_tot":len(list_models_to_gen) ,
                                       "nb_tasks_to_do":len(list_models_to_gen)*len(list_param)*nb_images_by_prompt ,
                                       "nb_tasks_tot":len(list_models_to_gen)*len(list_param)*nb_images_by_prompt,
@@ -366,21 +395,28 @@ def fonc_update_actu(text_actu,id):
     return gr.Textbox(s)
 
 def fonc_update_actu_2(id):
-    if id == 0:
+    if id == 0 :
+        return gr.Textbox("waiting...")
+    if cache_text_actu[str(id)]['warm']==True:
         return gr.Textbox("waiting...")
     s=""
-    s+=f"modules: {cache_text_actu[str(id)]['nb_modules_use']}/{nb_req_simult}\n"
-    s+=f"models remaining: {cache_text_actu[str(id)]['nb_models_to_do']}/{cache_text_actu[str(id)]['nb_models_tot']}\n"
     i=0
+    nb_ones=0
     for d in cache_text_actu[str(id)]['progress']:
         i+=1
+        if d==1:
+            nb_ones+=1
         s+=str(d)
         if i%10==0:
             s+=" "
         if i%50==0:
             s+="\n"
+    
     s+="\n"
-    s+=f"images remaining: {cache_text_actu[str(id)]['nb_tasks_to_do']}/{cache_text_actu[str(id)]['nb_tasks_tot']}\n"
+    s+=f"modules: {cache_text_actu[str(id)]['nb_modules_use']}/{nb_req_simult} ({cache_text_actu[str(id)]['nb_modules_use']-nb_ones}/{cache_text_actu[str(id)]['nb_modules_use']})\n"
+    s+=f"models remaining: {cache_text_actu[str(id)]['nb_models_to_do']}/{cache_text_actu[str(id)]['nb_models_tot']}\n"
+
+    s+=f"images remaining(done): {cache_text_actu[str(id)]['nb_tasks_to_do']}({cache_text_actu[str(id)]['nb_tasks_tot']-cache_text_actu[str(id)]['nb_tasks_to_do']})/{cache_text_actu[str(id)]['nb_tasks_tot']}\n"
     s+=f"fail attempt: {cache_text_actu[str(id)]['nb_fail']}\n"
     #s+=f"{tempo_update_actu*incr_update_actu[str(id)]} s"
     #incr_update_actu[str(id)]+=1
@@ -551,7 +587,103 @@ def fonc_add_to_text(text,list_models,index):
 
 def load_model_publ(choice_model_publ):
     return gr.Image(None,label=choice_model_publ,interactive=False),gr.Textbox(choice_model_publ,visible=False,show_label=False)
+
+
+def set_tasks_w(id_session,list_p,nb_i):
+    list_t=[]
+    for p in list_p:
+        cache_list_task_w[f"{id_session}"].append([nb_i,{"prompt":p[0],"nprompt":p[1],"width":p[2],"height":p[3],"steps":p[4],"cfg":p[5],"seed":p[6]},0])
+
+    cache_text_actu[f"{id_session}"]={"nb_modules_use":nb_req_simult,"stop":False,"nb_fail":0,"warm":True,"nb_prompt_to_do":len(list_p),"nb_prompt_tot":len(list_p),
+                                      "nb_tasks_to_do":len(list_p)*nb_i ,"nb_tasks_tot":len(list_p)*nb_i}
+    return 
+
+def fonc_start_w(id_session,id_module,s,cont,nb_modules,tag):
+    if cont==False or id_module>=nb_modules:
+        cache_text_actu[f"{id_session}"]["nb_modules_use"]-=1
+        print("manual stop")
+        return None,gr.Textbox(s),gr.Number(randint(1,MAX_SEED))
+
+    find_task=False
+    task={}
+    for t in cache_list_task_w[f"{id_session}"]:
+        if not find_task:
+            if t[0]>0:
+                t[0]-=1
+                t[2]+=1
+                task=t[1].copy()
+                find_task=True
+    if not find_task:
+        cache_text_actu[f"{id_session}"]["nb_modules_use"]-=1
+        return None,gr.Textbox(s),gr.Number(randint(1,MAX_SEED))
     
+    if tag == "":
+        tagT=["stable-diffusion-xl"]
+    else:
+        tagT=["stable-diffusion-xl",tag]
+    models_temp , models_plus_tags_temp = find_warm_model_list("John6666", tagT, "", "last_modified", 10000)
+    models_rand=[]
+    for m in models_temp:
+        if m in models:
+            models_rand.append(m)
+    if len(models_rand)!=0:
+        model_actu=models_temp[randint(0,len(models_rand)-1)]
+        print(f"find model : {model_actu}")
+    else:
+        print("no warm model")
+        cache_text_actu[f"{id_session}"]["nb_modules_use"]-=1
+        return None,gr.Textbox(s),gr.Number(randint(1,MAX_SEED))
+        
+        
+    print("begin gen image:")
+    print(model_actu)
+    print(task)
+    result=gen_fn(model_actu, task["prompt"], task["nprompt"], task["height"], task["width"], task["steps"], task["cfg"], task["seed"],inference_timeout_w)
+    print("reception")
+    if result!=None:
+        id_image=len(cache_image[f"{id_session}"])
+        for t in cache_list_task_w[f"{id_session}"]:
+            if t[1]==task:
+                t[2]-=1
+                if t[0]+t[2]<=0:
+                    cache_list_task_w[f"{id_session}"].remove(t)
+                    cache_text_actu[f"{id_session}"]["nb_prompt_to_do"]-=1 #################################
+                cache_text_actu[f"{id_session}"]["nb_tasks_to_do"]-=1
+                
+        task["id_image"]=id_image
+        task["model"]=model_actu
+        task["image"]=result
+        cache_image[f"{id_session}"].append(task)
+        print("image saved\n")
+        print(task)
+    else:
+        for t in cache_list_task_w[f"{id_session}"]:
+            if t[1]==task:
+                t[0]+=1
+                t[2]-=1
+        cache_text_actu[f"{id_session}"]["nb_fail"]+=1
+        print("fail to generate\n")
+    nb_task_to_do=cache_text_actu[str(id_session)]["nb_tasks_to_do"]
+    print(f"\n {nb_task_to_do} tasks to do\n")
+    return result , gr.Textbox(s+"1"),gr.Number(randint(1,MAX_SEED))
+
+
+def fonc_update_actu_w(id):
+    if id == 0 :
+        return gr.Textbox("waiting...")
+    if cache_text_actu[str(id)]['warm']==False:
+        return gr.Textbox("waiting...")
+    s=""
+    s+=f"modules: {cache_text_actu[str(id)]['nb_modules_use']}/{nb_req_simult} \n"
+    s+=f"prompts remaining: {cache_text_actu[str(id)]['nb_prompt_to_do']}/{cache_text_actu[str(id)]['nb_prompt_tot']}\n"
+
+    s+=f"images remaining(done): {cache_text_actu[str(id)]['nb_tasks_to_do']}({cache_text_actu[str(id)]['nb_tasks_tot']-cache_text_actu[str(id)]['nb_tasks_to_do']})/{cache_text_actu[str(id)]['nb_tasks_tot']}\n"
+    s+=f"fail attempt: {cache_text_actu[str(id)]['nb_fail']}\n"
+    #s+=f"{tempo_update_actu*incr_update_actu[str(id)]} s"
+    #incr_update_actu[str(id)]+=1
+    s+=f"{randint(1,MAX_SEED)}"
+    return gr.Textbox(s)
+
 def make_me():
         with gr.Tab(" "):
             with gr.Column():
@@ -638,6 +770,8 @@ def make_me():
                                          index_tag,index_first_model)
                         load_info=gr.Button("Load Models")
                         load_info.click(fonc_load_info,[nb_of_models_to_gen,index_tag,index_first_model],[nb_of_models_to_gen,disp_info,list_models_to_gen])
+                        button_load_random_models=gr.Button("Load Random Models")
+                        button_load_random_models.click(load_random_models,[nb_of_models_to_gen,index_tag],[nb_of_models_to_gen,disp_info,list_models_to_gen])
                     
                 with gr.Accordion("Models Custom",open=False) :
                     with gr.Row():
@@ -788,9 +922,134 @@ def make_me():
                             row_gallery.change(lambda g,h: gr.Gallery(g,rows=h),[gallery_by_prompt,row_gallery],[gallery_by_prompt])
                                 
     
-
             button_restore_session.click(fonc_restore_session,[id_session],[list_param,disp_param.dataset,list_models_to_gen,nb_of_models_to_gen])
 
+        with gr.Tab(" Warm ",visible=False) as tab_warm:  
+            button_test_pass.click(test_pass_aff,[test_pass],[tab_warm]) 
+            with gr.Column():
+                with gr.Group():
+                    with gr.Row():
+                        with gr.Column(scale=4):
+                            txt_input_w = gr.Textbox(label='Your prompt:', lines=4, interactive = True)
+                            neg_input_w = gr.Textbox(label='Negative prompt:', lines=4, interactive = True)
+                        with gr.Column(scale=4):
+                            with gr.Row():
+                                width_w = gr.Slider(label="Width", info="If 0, the default value is used.", maximum=2024, step=32, value=0, interactive = True)
+                                height_w = gr.Slider(label="Height", info="If 0, the default value is used.", maximum=2024, step=32, value=0, interactive = True)
+                            with gr.Row():
+                                choice_ratio_w = gr.Dropdown(label="Ratio Width/Height", 
+                                                    info="OverWrite Width and Height (W*H<1024*1024)", 
+                                                    show_label=True, choices=list(list_ratios) , interactive = True, value=list_ratios[0][1])
+                                choice_ratio_w.change(ratio_chosen,[choice_ratio_w,width,height_w],[width_w,height_w])
+                            with gr.Row():
+                                steps_w = gr.Slider(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=0, interactive = True)
+                                cfg_w = gr.Slider(label="Guidance scale", info="If 0, the default value is used.", maximum=30.0, step=0.1, value=0, interactive = True)
+                                seed_w = gr.Slider(label="Seed", info="Randomize Seed if -1.", minimum=-1, maximum=MAX_SEED, step=1, value=-1, interactive = True)
+                                
+                add_param_w=gr.Button("Add to the list")
+                del_param_w=gr.Button("Delete to the list")
+                
+                list_param_w=gr.Dropdown(choices=[["a",[["","",0,0,0,0,-1]]]], value=[["","",0,0,0,0,-1]], visible=False)
+                disp_param_w = gr.Examples(
+                    label="list of prompt",
+                    examples=list_param_w.value,
+                    inputs=[txt_input_w,neg_input_w,width_w,height_w,steps_w,cfg_w,seed_w],
+                    outputs=[txt_input_w,neg_input_w,width_w,height_w,steps_w,cfg_w,seed_w],
+                )
+                bool_warm=gr.Checkbox(visible=False,value=True)
+                add_param_w.click(fonc_add_param,[list_param_w,txt_input_w,neg_input_w,width_w,height_w,steps_w,cfg_w,seed_w],[disp_param_w.dataset,list_param_w])
+                add_param_w.click(set_session,[id_session,bool_warm],[id_session])
+                del_param_w.click(fonc_del_param,[list_param_w,txt_input_w,neg_input_w,width_w,height_w,steps_w,cfg_w,seed_w],[disp_param_w.dataset,list_param_w])
+                
+                with gr.Row():
+                    nb_images_by_prompt_w=gr.Number(2,label="Images by prompt",interactive = True)
+                    nb_modules_w=gr.Slider(label="Module use", minimum=1, maximum=nb_req_simult, step=1, value=40, interactive = True)
+                    text_tag_w=gr.Textbox("",label="Tag",interactive = True,lines=1)
+                
+                with gr.Column():
+
+                    button_load_task_w=gr.Button("LOAD")
+                    button_load_task_w.click(set_tasks_w,[id_session,list_param_w,nb_images_by_prompt_w],[])
+                    button_start_w=gr.Button("START")
+                    button_stop_w=gr.Button("STOP")
+                    cont_w=gr.Checkbox(True,visible=False)
+                    button_start_w.click(lambda:True,[],[cont_w])
+                    button_stop_w.click(lambda:False,[],[cont_w])
+                    text_actu_w=gr.Textbox(fonc_update_actu_w,inputs=id_session,every=tempo_update_actu,label="in progress",interactive=False,lines=6)############
+                    update_actu_w=gr.Number(0,visible=False)
+                    
+                    with gr.Accordion("Gallery Parameters",open=False) :
+                        with gr.Row():
+                            with gr.Column():
+                                set_height_gallery_w=gr.Checkbox(True,label="set height",show_label=True)
+                                height_gallery_w=gr.Number(650,label="height",show_label=True)
+                            col_gallery_w=gr.Number(5,label="nb columns",show_label=True)
+                            row_gallery_w=gr.Number(4,label="nb row",show_label=True)
+                            with gr.Column():
+                                button_reset_cache_image_w=gr.Button("Reset Images")
+                                button_reset_cache_image_w.click(reset_cache_image,[id_session],[])
+                                button_reset_cache_image_all_session_w=gr.Button("Reset Images ALL SESSION")
+                                button_reset_cache_image_all_session_w.click(reset_cache_image_all_sessions,[],[])
+                    
+                    with gr.Row():
+                        outputs_w=[]
+                        id_modules_w=[]
+                        states_w=[]
+                        for i in range(nb_req_simult):
+                            outputs_w.append(gr.Image(None,interactive=False,visible=False))
+                            id_modules_w.append(gr.Number(i,interactive=False,visible=False))
+                            states_w.append(gr.Textbox("1",interactive=False,visible=False))
+                        for o,i,s in zip(outputs_w,id_modules_w,states_w):
+                            s.change(fonc_start_w,[id_session,i,s,cont_w,nb_modules_w,text_tag_w],[o,s,update_actu_w])#################
+                            gen_event = gr.on(triggers=[button_start_w.click], fn=fonc_init,inputs=[s], outputs=[s])###################
+
+                            
+                        with gr.Column(scale=2):    
+                            gallery_w = gr.Gallery(label="Output", show_download_button=True, elem_classes="gallery",
+                                            interactive=False, show_share_button=False, container=True, format="png",
+                                            preview=True, object_fit="contain",columns=5,rows=4,height=650)
+                            gallery_by_prompt_w = gr.Gallery(label="Output", show_download_button=True, elem_classes="gallery",
+                                            interactive=False, show_share_button=False, container=True, format="png",
+                                            preview=True, object_fit="contain",columns=5,rows=4,visible=False,height=650)
+                            
+                            set_height_gallery_w.change(lambda g,h,s: gr.Gallery(g,height=h) if s else gr.Gallery(g,height=None),
+                                                      [gallery_w,height_gallery_w,set_height_gallery_w],[gallery_w])
+                            height_gallery_w.change(lambda g,h: gr.Gallery(g,height=h),[gallery_w,height_gallery_w],[gallery_w])
+                            col_gallery_w.change(lambda g,h: gr.Gallery(g,columns=h),[gallery_w,col_gallery_w],[gallery_w])
+                            row_gallery_w.change(lambda g,h: gr.Gallery(g,rows=h),[gallery_w,row_gallery_w],[gallery_w])
+                            
+                            set_height_gallery_w.change(lambda g,h,s: gr.Gallery(g,height=h) if s else gr.Gallery(g,height=None),
+                                                      [gallery_by_prompt_w,height_gallery_w,set_height_gallery_w],[gallery_by_prompt_w])
+                            height_gallery_w.change(lambda g,h: gr.Gallery(g,height=h),[gallery_by_prompt_w,height_gallery_w],[gallery_by_prompt_w])
+                            col_gallery_w.change(lambda g,h: gr.Gallery(g,columns=h),[gallery_by_prompt_w,col_gallery_w],[gallery_by_prompt_w])
+                            row_gallery_w.change(lambda g,h: gr.Gallery(g,rows=h),[gallery_by_prompt_w,row_gallery_w],[gallery_by_prompt_w])
+                            
+                        with gr.Column(scale=3):
+                            button_load_gallery_w=gr.Button("Load Gallery All")
+                            button_load_gallery_w.click(fonc_load_gallery,[id_session,gallery_w],[gallery_w])
+                            
+                            with gr.Accordion("Gallery by Prompt",open=False) :
+                                index_gallery_by_prompt_w=gr.Number(0,visible=False)
+                                button_load_gallery_by_prompt_w=gr.Button("Load Gallery by prompt")
+                                text_gallery_by_prompt_w=gr.Textbox(f"{index_gallery_by_prompt_w.value+1}/{len(list_param_w.value)}",show_label=False)
+                                index_gallery_by_prompt_w.change(lambda i,p:gr.Textbox(f"{i+1}/{len(p)}"),[index_gallery_by_prompt_w,list_param_w],[text_gallery_by_prompt_w])
+                                button_load_gallery_by_prompt_w.click(load_gallery_by_prompt,
+                                                                    [id_session,gallery_by_prompt_w,index_gallery_by_prompt_w,list_param_w],[gallery_by_prompt_w])
+                                gen_event_gallery_by_prompt_w = gr.on(triggers=[button_load_gallery_by_prompt_w.click], fn=lambda g:gr.Gallery(g,visible=False),
+                                                                inputs=[gallery_w], outputs=[gallery_w])
+                                gen_event_gallery_w = gr.on(triggers=[button_load_gallery_w.click], fn=lambda g:gr.Gallery(g,visible=False),
+                                                                inputs=[gallery_by_prompt_w], outputs=[gallery_by_prompt_w])
+                                with gr.Row():
+                                    button_gallery_prev_prompt_w=gr.Button("Prev prompt")
+                                    button_gallery_next_prompt_w=gr.Button("Next prompt")
+                                    button_gallery_next_prompt_w.click(lambda i,p: (i+1)%len(p),[index_gallery_by_prompt_w,list_param_w],[index_gallery_by_prompt_w])
+                                    button_gallery_prev_prompt_w.click(lambda i,p: (i-1)%len(p),[index_gallery_by_prompt_w,list_param_w],[index_gallery_by_prompt_w])
+                                index_gallery_by_prompt_w.change(load_gallery_by_prompt,
+                                                                    [id_session,gallery_by_prompt_w,index_gallery_by_prompt_w,list_param_w],[gallery_by_prompt_w])
+                        
+                    
+
+