Update app.py

app.py

@@ -461,154 +461,109 @@ pipe4.fuse_lora()
 # pipe3.scheduler.set_timesteps(50)
 ######### Gradio Fuction #############
 
-def swap_to_gallery(images):
-    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
-
-def upload_example_to_gallery(images, prompt, style, negative_prompt):
-    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
-
-def remove_back_to_files():
-    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
-
 def remove_tips():
     return gr.update(visible=False)
 
 def apply_style_positive(style_name: str, positive: str):
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive) 
+    return p.replace("{prompt}", positive)
 
 def apply_style(style_name: str, positives: list, negative: str = ""):
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return [p.replace("{prompt}", positive) for positive in positives], n + ' ' + negative
 
-def change_visiale_by_model_type(_model_type):
-    if _model_type == "Only Using Textual Description":
-        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-    elif _model_type == "Using Ref Images":
-        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
-    else:
-        raise ValueError("Invalid model type",_model_type)
+def change_visible_by_model_type(_model_type):
+    # Since you are **only using text**, always hide ref image uploads
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 
 ######### Image Generation ##############
 @spaces.GPU(duration=120)
-def process_generation(_sd_type,_model_type,_upload_images, _num_steps,style_name, _Ip_Adapter_Strength ,_style_strength_ratio, guidance_scale, seed_,  sa32_, sa64_, id_length_,  general_prompt, negative_prompt,prompt_array,G_height,G_width,_comic_type):
-    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
-    if _model_type == "Photomaker" and "img" not in general_prompt:
-        raise gr.Error("Please add the triger word \" img \"  behind the class word you want to customize, such as: man img or woman img")
-    if _upload_images is None and _model_type != "original":
-        raise gr.Error(f"Cannot find any input face image!")
-    if len(prompt_array.splitlines()) > 10:
-        raise gr.Error(f"No more than 10 prompts in huggface demo for Speed! But found {len(prompt_array.splitlines())} prompts!")
-    global sa32, sa64,id_length,total_length,attn_procs,unet,cur_model_type,device
+def process_generation(_sd_type, _num_steps, style_name, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, G_height, G_width, _comic_type):
+    global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type, device
     global num_steps
     global write
-    global cur_step,attn_count
-    global height,width
+    global cur_step, attn_count
+    global height, width
     height = G_height
     width = G_width
-    global pipe2,pipe4
-    global sd_model_path,models_dict
+    global pipe2, pipe4
+    global sd_model_path, models_dict
+
     sd_model_path = models_dict[_sd_type]
-    num_steps =_num_steps
+    num_steps = _num_steps
     use_safe_tensor = True
-    if  style_name == "(No style)":
+
+    if style_name == "(No style)":
         sd_model_path = models_dict["SDXL"]
-    if _model_type == "original":
-        pipe = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16)
-        pipe = pipe.to(device)
-        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
-        # pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-        # pipe.scheduler.set_timesteps(50)
-        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
-    elif _model_type == "Photomaker":
-        if _sd_type != "SDXL" and style_name != "(No style)":
-            pipe = pipe2.to(device)
-            pipe.id_encoder.to(device)
-            set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
-        else:
-            pipe = pipe4.to(device)
-            pipe.id_encoder.to(device)
-            set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
-    else:
-        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
-        ##### ########################
-    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    pipe = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16)
+    pipe = pipe.to(device)
     pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
-    cur_model_type = _sd_type+"-"+_model_type+""+str(id_length_)
-    if _model_type != "original":
-        input_id_images = []
-        for img in _upload_images:
-            print(img)
-            input_id_images.append(load_image(img))
+    set_attention_processor(pipe.unet, id_length_, is_ipadapter=False)
+
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    cur_model_type = _sd_type + "-original-" + str(id_length_)
+    
     prompts = prompt_array.splitlines()
-    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
-    if start_merge_step > 30:
-        start_merge_step = 30
-    print(f"start_merge_step:{start_merge_step}")
+    if len(prompts) > 10:
+        raise gr.Error(f"No more than 10 prompts in the Hugging Face demo for speed! But found {len(prompts)} prompts!")
+
     generator = torch.Generator(device="cuda").manual_seed(seed_)
-    sa32, sa64 =  sa32_, sa64_
+    sa32, sa64 = sa32_, sa64_
     id_length = id_length_
     clipped_prompts = prompts[:]
-    prompts = [general_prompt + "," + prompt if "[NC]" not in prompt else prompt.replace("[NC]","")  for prompt in clipped_prompts]
+
+    prompts = [general_prompt + "," + prompt if "[NC]" not in prompt else prompt.replace("[NC]", "") for prompt in clipped_prompts]
     prompts = [prompt.rpartition('#')[0] if "#" in prompt else prompt for prompt in prompts]
-    print(prompts)
+
     id_prompts = prompts[:id_length]
     real_prompts = prompts[id_length:]
+
     torch.cuda.empty_cache()
     write = True
     cur_step = 0
-
     attn_count = 0
+
     id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+
     setup_seed(seed_)
+
     total_results = []
-    if _model_type == "original":
-        id_images = pipe(id_prompts, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
-    elif _model_type == "Photomaker":
-        id_images = pipe(id_prompts,input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale, start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
-    else: 
-        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+    
+    # Generate ID images
+    id_images = pipe(id_prompts, num_inference_steps=num_steps, guidance_scale=guidance_scale,
+                     height=height, width=width, negative_prompt=negative_prompt, generator=generator).images
     total_results = id_images + total_results
     yield total_results
+
+    # Generate real comic images
     real_images = []
     write = False
     for real_prompt in real_prompts:
         setup_seed(seed_)
         cur_step = 0
         real_prompt = apply_style_positive(style_name, real_prompt)
-        if _model_type == "original":   
-            real_images.append(pipe(real_prompt,  num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
-        elif _model_type == "Photomaker":      
-            real_images.append(pipe(real_prompt, input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
-        else:
-            raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+        
+        real_images.append(pipe(real_prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale,
+                                height=height, width=width, negative_prompt=negative_prompt, generator=generator).images[0])
         total_results = [real_images[-1]] + total_results
         yield total_results
+
+    # Comic typesetting if selected
     if _comic_type != "No typesetting (default)":
-        captions= prompt_array.splitlines()
-        captions = [caption.replace("[NC]","") for caption in captions]
+        captions = prompt_array.splitlines()
+        captions = [caption.replace("[NC]", "") for caption in captions]
         captions = [caption.split('#')[-1] if "#" in caption else caption for caption in captions]
         from PIL import ImageFont
-        total_results = get_comic(id_images + real_images, _comic_type,captions= captions,font=ImageFont.truetype("./Inkfree.ttf", int(45))) + total_results
-    if _model_type == "Photomaker":
-        pipe = pipe2.to("cpu")
-        pipe.id_encoder.to("cpu")
-        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
-    yield total_results
-
+        total_results = get_comic(id_images + real_images, _comic_type, captions=captions,
+                                  font=ImageFont.truetype("./Inkfree.ttf", int(45))) + total_results
 
+    yield total_results
 
 def array2string(arr):
-    stringtmp = ""
-    for i,part in enumerate(arr):
-        if i != len(arr)-1:
-            stringtmp += part +"\n"
-        else:
-            stringtmp += part
-
-    return stringtmp
-
+    return "\n".join(arr)
 
 #################################################
 #################################################
@@ -617,167 +572,62 @@ with gr.Blocks(css=css) as demo:
     binary_matrixes = gr.State([])
     color_layout = gr.State([])
 
-    # gr.Markdown(logo)
     gr.Markdown(title)
     gr.Markdown(description)
 
     with gr.Row():
         with gr.Group(elem_id="main-image"):
-            # button_run = gr.Button("generate id images ! 😺", elem_id="main_button", interactive=True)
-
-            prompts = []
-            colors = []
-            # with gr.Column(visible=False) as post_sketch:
-            #     for n in range(MAX_COLORS):
-            #         if n == 0 :
-            #             with gr.Row(visible=False) as color_row[n]:
-            #                 colors.append(gr.Image(shape=(100, 100), label="background", type="pil", image_mode="RGB", width=100, height=100))
-            #                 prompts.append(gr.Textbox(label="Prompt for the background (white region)", value=""))
-            #         else:
-            #             with gr.Row(visible=False) as color_row[n]:
-            #                 colors.append(gr.Image(shape=(100, 100), label="segment "+str(n), type="pil", image_mode="RGB", width=100, height=100))
-            #                 prompts.append(gr.Textbox(label="Prompt for the segment "+str(n)))
-
-            #     get_genprompt_run = gr.Button("(2) I've finished segment labeling ! 😺", elem_id="prompt_button", interactive=True)
-
-            with gr.Column(visible=True) as gen_prompt_vis:
-                sd_type = gr.Dropdown(choices=list(models_dict.keys()), value = "SDXL",label="sd_type", info="Select pretrained model")
-                model_type = gr.Radio(["Only Using Textual Description", "Using Ref Images"], label="model_type", value = "Only Using Textual Description",  info="Control type of the Character")
-                with gr.Group(visible=False) as control_image_input:
-                    files = gr.Files(
-                                label="Drag (Select) 1 or more photos of your face",
-                                file_types=["image"],
-                            )
-                    uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
-                    with gr.Column(visible=False) as clear_button:
-                        remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+
+            with gr.Column(visible=True):
+                sd_type = gr.Dropdown(choices=list(models_dict.keys()), value="SDXL", label="sd_type", info="Select pretrained model")
+
                 general_prompt = gr.Textbox(value='', label="(1) Textual Description for Character", interactive=True)
-                negative_prompt = gr.Textbox(value='', label="(2) Negative_prompt", interactive=True)
-                style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
-                prompt_array = gr.Textbox(lines = 3,value='', label="(3) Comic Description (each line corresponds to a frame).", interactive=True)
-                with gr.Accordion("(4) Tune the hyperparameters", open=True):
-                    #sa16_ = gr.Slider(label=" (The degree of Paired Attention at 16 x 16 self-attention layers) ", minimum=0, maximum=1., value=0.3, step=0.1)
-                    sa32_ = gr.Slider(label=" (The degree of Paired Attention at 32 x 32 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
-                    sa64_ = gr.Slider(label=" (The degree of Paired Attention at 64 x 64 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
-                    id_length_ = gr.Slider(label= "Number of id images in total images" , minimum=2, maximum=4, value=3, step=1)
-                    # total_length_ = gr.Slider(label= "Number of total images", minimum=1, maximum=20, value=1, step=1)
+                negative_prompt = gr.Textbox(value='', label="(2) Negative Prompt", interactive=True)
+                style = gr.Dropdown(label="Style Template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                prompt_array = gr.Textbox(lines=3, value='', label="(3) Comic Description (each line = one frame)", interactive=True)
+
+                with gr.Accordion("(4) Tune the Hyperparameters", open=True):
+                    sa32_ = gr.Slider(label="Paired Attention at 32x32 layers", minimum=0, maximum=1., value=0.7, step=0.1)
+                    sa64_ = gr.Slider(label="Paired Attention at 64x64 layers", minimum=0, maximum=1., value=0.7, step=0.1)
+                    id_length_ = gr.Slider(label="Number of id images", minimum=2, maximum=4, value=3, step=1)
                     seed_ = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, value=0, step=1)
-                    num_steps = gr.Slider( 
-                        label="Number of sample steps",
-                        minimum=25,
-                        maximum=50,
-                        step=1,
-                        value=50,
-                    )
-                    G_height = gr.Slider( 
-                        label="height",
-                        minimum=256,
-                        maximum=1024,
-                        step=32,
-                        value=1024,
-                    )
-                    G_width = gr.Slider( 
-                        label="width",
-                        minimum=256,
-                        maximum=1024,
-                        step=32,
-                        value=1024,
-                    )
-                    comic_type = gr.Radio(["No typesetting (default)", "Four Pannel", "Classic Comic Style"], value = "Classic Comic Style", label="Typesetting Style", info="Select the typesetting style ")
-                    guidance_scale = gr.Slider(
-                        label="Guidance scale",
-                        minimum=0.1,
-                        maximum=10.0,
-                        step=0.1,
-                        value=5,
-                    )
-                    style_strength_ratio = gr.Slider(
-                        label="Style strength of Ref Image (%)",
-                        minimum=15,
-                        maximum=50,
-                        step=1,
-                        value=20,
-                        visible=False
-                    )
-                    Ip_Adapter_Strength = gr.Slider(
-                        label="Ip_Adapter_Strength",
-                        minimum=0,
-                        maximum=1,
-                        step=0.1,
-                        value=0.5,
-                        visible=False
-                    )
-                final_run_btn = gr.Button("Generate ! 😺")
+                    num_steps = gr.Slider(label="Number of Sample Steps", minimum=25, maximum=50, step=1, value=50)
+                    G_height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=1024)
+                    G_width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=1024)
+                    comic_type = gr.Radio(["No Typesetting (default)", "Four Panel", "Classic Comic Style"], value="Classic Comic Style", label="Typesetting Style")
+                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=10.0, step=0.1, value=5)
 
+                final_run_btn = gr.Button("Generate ! 😺")
 
         with gr.Column():
             out_image = gr.Gallery(label="Result", columns=2, height='auto')
-            generated_information = gr.Markdown(label="Generation Details", value="",visible=False)
+            generated_information = gr.Markdown(label="Generation Details", value="", visible=False)
             gr.Markdown(version)
-    model_type.change(fn = change_visiale_by_model_type , inputs = model_type, outputs=[control_image_input,style_strength_ratio,Ip_Adapter_Strength])
-    files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
-    remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
-
-    final_run_btn.click(fn=set_text_unfinished, outputs = generated_information
-    ).then(process_generation, inputs=[sd_type,model_type,files, num_steps,style, Ip_Adapter_Strength,style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array,G_height,G_width,comic_type], outputs=out_image
-    ).then(fn=set_text_finished,outputs = generated_information)
 
+    final_run_btn.click(fn=set_text_unfinished, outputs=generated_information
+    ).then(
+        process_generation,
+        inputs=[sd_type, num_steps, style, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, G_height, G_width, comic_type],
+        outputs=out_image
+    ).then(fn=set_text_finished, outputs=generated_information)
 
     gr.Examples(
         examples=[
-            [0,0.5,0.5,2,"a scientist, wearing a lab coat and glasses",
-                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
-                array2string(["in the lab, surrounded by bubbling beakers #Today I will discover something new",
-                            "[NC]He discovers a glowing chemical compound #Wow!!!",
-                            "[NC] The car on the road, near the forest #He drives to the forest in search of treasure.",
-                            "startled, holding the beaker #The compound explodes with a bright light",
-                            "[NC]He opens his eyes to find himself in a different world"
-                            ]),
-                            "Comic book","Only Using Textual Description",get_image_path('./yannlecun.png'),768,768
-            ],
-            [0,0.5,0.5,2,"a cyberpunk hacker, glowing wires, neon glasses",
-                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
-                array2string(["[NC]In a dark room filled with monitors #She types rapidly on a neon-lit keyboard",
-                            "neon city street at night, people walking by",
-                            "a robot chases her through a back alley",
-                            "[NC]She jumps onto a rooftop, escaping",
-                            ]),
-                            "Japanese Anime","Using Ref Images",get_image_path('./ironman.jpeg'),768,768
-            ],
-            [1,0.5,0.5,3,"a knight in silver armor with blue cape",
-                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
-                   array2string(["[NC]On a grassy battlefield #He raises his sword toward the enemy",
-                                "charging forward on horseback",
-                                "blocking an arrow with his shield",
-                                "[NC]He wins the battle and is crowned by the king"]),
-                                "Japanese Anime",  "Using Ref Images",get_image_path('./taylor.png'),768,768
-                ],
-                [0,0.5,0.5,3,"a girl with pink hair, school uniform",
-                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
-                   array2string(["at the school gate",
-                                "[NC]She notices a mysterious envelope in her locker",
-                                "reading the letter in shock #It's a confession of love",
-                                "running outside, blushing"
-                                ]),
-                                "Japanese Anime","Only Using Textual Description",get_image_path('./yannlecun.png'),768,768
-                ],
-                [0,0.3,0.5,3,"an astronaut in white spacesuit",
-                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
-                    array2string([
-                            "floating in space above Earth",
-                            "[NC]Spots a mysterious alien ship in the distance",
-                            "enters the ship cautiously",
-                            "[NC]finds a message written in glowing symbols"]),
-                    "Japanese Anime","Only Using Textual Description",get_image_path('./ironman.jpeg'),768,768
-                ]
-                ],
-        inputs=[seed_, sa32_, sa64_, id_length_,  general_prompt, negative_prompt, prompt_array,style,model_type,files,G_height,G_width],
-        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
-        # run_on_click=True,
+            [0, 0.5, 0.5, 2, "a scientist, wearing a lab coat and glasses",
+             "bad anatomy, bad hands, missing fingers, ...",
+             array2string([
+                 "In the lab, surrounded by bubbling beakers #Today I will discover something new",
+                 "[NC]He discovers a glowing chemical compound #Wow!!!",
+                 "[NC]The car on the road, near the forest #He drives to the forest in search of treasure.",
+                 "startled, holding the beaker #The compound explodes with a bright light",
+                 "[NC]He opens his eyes to find himself in a different world"
+             ]),
+             "Comic book", 768, 768],
+        ],
+        inputs=[seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, style, G_height, G_width],
         label='😺 Examples 😺',
     )
-    gr.Markdown(article)
 
-    # demo.load(None, None, None, _js=load_js)
+    gr.Markdown(article)
 
 demo.launch()