Update app.py

app.py

@@ -39,7 +39,7 @@ pipe = ConsistentIDStableDiffusionPipeline.from_pretrained(
 ## BiSenet
 bise_net_cp_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="face_parsing.pth", local_dir="./checkpoints")
 bise_net = BiSeNet(n_classes = 19)
-bise_net.load_state_dict(torch.load(bise_net_cp_path, map_location="cpu"))
+bise_net.load_state_dict(torch.load(bise_net_cp_path, map_location="cpu")).to(device)
 
 ### Load consistentID_model checkpoint
 pipe.load_ConsistentID_model(
@@ -52,6 +52,27 @@ pipe.load_ConsistentID_model(
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
 
+# @torch.inference_mode()
+# def Enhance_prompt(prompt,select_images):
+    
+#     llva_prompt = f'Please ignore the image. Enhance the following text prompt for me. You can associate more details with the character\'s gesture, environment, and decent clothing:"{prompt}".' 
+#     args = type('Args', (), {
+#         "model_path": llva_model_path,
+#         "model_base": None,
+#         "model_name": get_model_name_from_path(llva_model_path),
+#         "query": llva_prompt,
+#         "conv_mode": None,
+#         "image_file": select_images,
+#         "sep": ",",
+#         "temperature": 0,
+#         "top_p": None,
+#         "num_beams": 1,
+#         "max_new_tokens": 512
+#     })() 
+#     Enhanced_prompt = eval_model(args, llva_tokenizer, llva_model, llva_image_processor)
+
+#     return Enhanced_prompt
+    
 @spaces.GPU
 def process(inputImage,prompt,negative_prompt):
 
@@ -65,14 +86,15 @@ def process(inputImage,prompt,negative_prompt):
         prompt = "A man, with backpack, in a raining tropical forest, adventuring, holding a flashlight, in mist, seeking animals"
         prompt = "A person, in a sowm, wearing santa hat and a scarf, with a cottage behind"
     else:
-        prompt=Enhance_prompt(prompt,blank_image)
+        # prompt=Enhance_prompt(prompt,blank_image) # TODO
+        prompt=prompt
         print(prompt)
         pass
         
     if negative_prompt == "":
         negative_prompt = ",monochrome, lowres, bad anatomy, worst quality, low quality, blurry"
 
-    #Extend Prompt
+    # Extend Prompt
     prompt = "cinematic photo," + prompt + ", 50mm photograph, half-length portrait, film, bokeh, professional, 4k, highly detailed"
 
     negtive_prompt_group="((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))). out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))"
@@ -103,21 +125,11 @@ def process(inputImage,prompt,negative_prompt):
 
     return os.path.join(output_dir, f"{current_date}-{seed}.jpg")
 
-
-iface = gr.Interface(
-    fn=process,
-    inputs=[
-        gr.Image(label="Upload Image"), 
-        gr.Textbox(label="prompt",placeholder="A man, in a forest, adventuring"),
-        gr.Textbox(label="negative prompt",placeholder="monochrome, lowres, bad anatomy, worst quality, low quality, blurry"),
-    ],
-    outputs=[
-        gr.Image(label="Output"), 
-    ],
-    title="ConsistentID Demo",
-    description="Put reference portrait below" ,
-    allow_flagging="never"
-)
+# # Gets the templates
+# script_directory = os.path.dirname(os.path.realpath(__file__))
+# # preset_template = glob.glob(script_directory+"/images/templates/*.png")
+# preset_template = glob.glob("./images/templates/*.png")
+# preset_template = preset_template + glob.glob("./images/templates/*.jpg")
 
 # # Use Blocks Create Gradio 
 # with gr.Blocks(title="ConsistentID Demo") as demo:
@@ -185,6 +197,21 @@ iface = gr.Interface(
 #         btn.click(fn=process, inputs=[selected_template_images,costum_image,prompt,nagetive_prompt,prompt_selected,retouching
 #             ,model_selected_tab,prompt_selected_tab,width,height,merge_steps], outputs=out)
 
+iface = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Image(label="Upload Image"), 
+        gr.Textbox(label="prompt",placeholder="A man, in a forest, adventuring"),
+        gr.Textbox(label="negative prompt",placeholder="monochrome, lowres, bad anatomy, worst quality, low quality, blurry"),
+    ],
+    outputs=[
+        gr.Image(label="Output"), 
+    ],
+    title="ConsistentID Demo",
+    description="Put reference portrait below" ,
+    allow_flagging="never"
+)
+
 iface.launch() # zero.device
 
 # @spaces.GPU