Update app.py

app.py

@@ -106,6 +106,8 @@ def select_segment(segment_output, segment_name):
 
 def return_image(imageEditor):
     return imageEditor['composite']
+def return_image2(image):
+    return image
 
 def rembg_client(request: gr.Request):
     try:
@@ -190,9 +192,26 @@ def upscale_image(image_pil, version="v1.4", rescaling_factor=None):
 
 #     return inpainting_image(imageEditor['composite'])
 
-def inpainting_run(model_name, use_rasg, use_painta, prompt, imageMask,
-    hr_image, seed, eta, negative_prompt, positive_prompt, ddim_steps,
-    guidance_scale=7.5, batch_size=1, session_id=''
+
+def inpainting_client(request: gr.Request):
+    try:
+        client = Client("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU", headers={"X-IP-Token": request.headers['x-ip-token']})
+        print("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU Ip token")
+        return client
+    except:
+        print("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU no token")
+        return Client("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU")
+
+def inpainting_run(input_image_editor,
+            prompt,
+            negative_prompt,
+            controlnet_conditioning_scale,
+            guidance_scale,
+            seed,
+            num_inference_steps,
+            true_guidance_scale,
+            
+            request: gr.Request
 ):
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
         # hr_image.save(temp_file.name)
@@ -205,20 +224,21 @@ def inpainting_run(model_name, use_rasg, use_painta, prompt, imageMask,
             print("маска", temp_file_path2)
 
 
-            client = Client("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU")
+            # client = Client("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU")
+            client = inpainting_client(request)
             result = client.predict(
                 input_image_editor={"background":handle_file(temp_file_path),"layers":[handle_file(temp_file_path2)],"composite":None},
                 prompt=prompt,
                 negative_prompt=negative_prompt,
-                controlnet_conditioning_scale=0.9,
-                guidance_scale=3.5,
-                seed=124,
-                num_inference_steps=24,
-                true_guidance_scale=3.5,
+                controlnet_conditioning_scale=controlnet_conditioning_scale,
+                guidance_scale=guidance_scale,
+                seed=seed,
+                num_inference_steps=num_inference_steps,
+                true_guidance_scale=true_guidance_scale,
                 api_name="/process"
             )
             print(result)
-            return [(result, 'label')]
+            return result
 
 
             # client = Client("https://pair-hd-painter.hf.space/--replicas/0l7ng/")
@@ -311,12 +331,6 @@ def generate_3d_model2(image_pil, request: gr.Request):
 
 ### some configs
 
-inpainting_models = OrderedDict([
-    ("Dreamshaper Inpainting V8", 'ds8_inp'),
-    ("Stable-Inpainting 2.0", 'sd2_inp'),
-    ("Stable-Inpainting 1.5", 'sd15_inp')
-])
-
 negative_prompt_str = "text, bad anatomy, bad proportions, blurry, cropped, deformed, disfigured, duplicate, error, extra limbs, gross proportions, jpeg artifacts, long neck, low quality, lowres, malformed, morbid, mutated, mutilated, out of frame, ugly, worst quality"
 positive_prompt_str = "Full HD, 4K, high quality, high resolution"
 
@@ -338,6 +352,7 @@ with gr.Blocks() as demo:
                 with gr.Row(equal_height=True):
                     segment_text_input = gr.Textbox(label="Имя сегмента для дальнейшего редактирования")
                     select_segment_button = gr.Button("Использовать сегмент")
+
     with gr.Tab("Редактирование"):
         with gr.Row(equal_height=True):
             with gr.Column(scale=5):
@@ -375,85 +390,32 @@ with gr.Blocks() as demo:
                     add_transparent_border_button = gr.Button("Добавить прозрачные края")
                 use_inpainting_button = gr.Button("Использовать сегмент для Inpainting")
                 use_button = gr.Button("Использовать сегмент для 3D")
+
     with gr.Tab("Inpainting"):
-        # with gr.Row(equal_height=True):
-        #     with gr.Column(scale=5):
-        #         inpainting_input = gr.ImageEditor(type="pil", label="Сегмент для Inpainting")
-        #     with gr.Column(scale=5):
-        #         inpainting_text_input = gr.Textbox(label="Текст для Inpainting")
-        #         inpainting_button = gr.Button("Inpainting")
-        #         use_inpainting_button2 = gr.Button("Вернуться к редактированию")
-        #         use_button2 = gr.Button("Использовать сегмент для 3D")
-
-        with gr.Row():
-            session_id = gr.Textbox(value=str(uuid.uuid4()), visible=False)
-            with gr.Column():
-                model_picker = gr.Dropdown(
-                    list(inpainting_models.keys()),
-                    value=list(inpainting_models.keys())[0],
-                    label = "Please select a model!",
-                )
-            with gr.Column():
-                use_painta = gr.Checkbox(value = True, label = "Use PAIntA")
-                use_rasg = gr.Checkbox(value = True, label = "Use RASG")
-        prompt = gr.Textbox(label = "Inpainting Prompt")
-        with gr.Row():
-            with gr.Column():
-                # imageMask = gr.ImageMask(label = "Input Image", brush_color='#ff0000', elem_id="inputmask", type="pil")
-                # imageMask = gr.ImageMask(label = "Input Image", elem_id="inputmask", type="pil")
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=5):
+                # inpainting_input = gr.ImageEditor(type="pil", label="Сегмент для Inpainting")
                 imageMask = gr.ImageEditor(
-                    label='Image',
+                    label='Сегмент для Inpainting',
                     type='pil',
                     sources=["upload", "webcam"],
                     image_mode='RGB',
                     layers=False,
                     brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"))
-                hr_image = gr.Image(visible=False, type="pil")
-
-                with gr.Row():
-                    inpaint_btn = gr.Button("Inpaint", scale = 0)
-      
+                
+                prompt = gr.Textbox(lines=2, label="Введите промпт для Inpainting", placeholder="Enter prompt here...")
+                inpainting_button = gr.Button("Inpainting")
                 with gr.Accordion('Advanced options', open=False):
-                    guidance_scale = gr.Slider(minimum = 0, maximum = 30, value = 7.5, label = "Guidance Scale")
-                    eta = gr.Slider(minimum = 0, maximum = 1, value = 0.1, label = "eta")
-                    ddim_steps = gr.Slider(minimum = 10, maximum = 100, value = 50, step =  1, label = 'Number of diffusion steps')
-                    with gr.Row():
-                        seed = gr.Number(value = 49123, label = "Seed")
-                        batch_size = gr.Number(value = 1, label = "Batch size", minimum=1, maximum=4) 
-                    negative_prompt = gr.Textbox(value=negative_prompt_str, label = "Negative prompt", lines=3)
-                    positive_prompt = gr.Textbox(value=positive_prompt_str, label = "Positive prompt", lines=1)
-
-            with gr.Column():
-                with gr.Row():
-                    output_gallery = gr.Gallery(
-                        [],
-                        columns = 4,
-                        preview = True,
-                        allow_preview = True,
-                        object_fit='scale-down',
-                        elem_id='outputgallery'
-                    )
-        inpaint_btn.click(
-            fn=inpainting_run, 
-            inputs=[
-                model_picker,
-                use_rasg,
-                use_painta,
-                prompt,
-                imageMask,
-                hr_image,
-                seed,
-                eta,
-                negative_prompt,
-                positive_prompt,
-                ddim_steps,
-                guidance_scale,
-                batch_size,
-                session_id
-            ], 
-            outputs=output_gallery, 
-            api_name="inpaint"
-        )  
+                    negative_prompt = gr.Textbox(lines=2, value=negative_prompt_str, label="Negative prompt", placeholder="Enter negative_prompt here...")
+                    controlnet_conditioning_scale = gr.Slider(minimum=0, step=0.01, maximum=1, value=0.9, label="controlnet_conditioning_scale")
+                    guidance_scale = gr.Slider(minimum=1, step=0.5, maximum=10, value=3.5, label="Image to generate")
+                    seed  = gr.Slider(minimum=0, step=1, maximum=10000000, value=124, label="Seed Value")
+                    num_inference_steps = gr.Slider(minimum=1, step=1, maximum=30, value=24, label="num_inference_steps")
+                    true_guidance_scale = gr.Slider(minimum=1, step=1, maximum=10, value=3.5, label="true_guidance_scale")
+            with gr.Column(scale=5):
+                after_inpainting = gr.Image(type="pil", label="Изображение после Inpainting")
+                use_inpainting_button2 = gr.Button("Вернуться к редактированию")
+                use_button2 = gr.Button("Использовать сегмент для 3D")
 
     with gr.Tab("Создание 3D"):
         with gr.Row(equal_height=True):
@@ -476,13 +438,27 @@ with gr.Blocks() as demo:
     rembg_button.click(rembg, inputs=[segment_input, rembg_model_selector], outputs=segment_input)
     remove_background_button.click(remove_black_make_transparent, inputs=segment_input, outputs=segment_input)
     add_transparent_border_button.click(add_transparent_border, inputs=[segment_input, add_transparent_border_slider], outputs=segment_input)
-    # use_inpainting_button.click(return_image, inputs=segment_input, outputs=inpainting_input)
+    use_inpainting_button.click(return_image, inputs=segment_input, outputs=imageMask)
     use_button.click(return_image, inputs=segment_input, outputs=segment_3d_input)
 
     #tab3
     # inpainting_button.click(inpainting, inputs=inpainting_input, outputs=inpainting_input)
-    # use_inpainting_button2.click(return_image, inputs=inpainting_input, outputs=segment_input)
-    # use_button2.click(return_image, inputs=inpainting_input, outputs=segment_3d_input)
+    inpainting_button.click(
+        fn=inpainting_run,
+        inputs=[
+            imageMask,
+            prompt,
+            negative_prompt,
+            controlnet_conditioning_scale,
+            guidance_scale,
+            seed,
+            num_inference_steps,
+            true_guidance_scale
+        ],
+        outputs=after_inpainting
+    )
+    use_inpainting_button2.click(return_image2, inputs=after_inpainting, outputs=segment_input)
+    use_button2.click(return_image2, inputs=after_inpainting, outputs=segment_3d_input)
 
     #3d buttons
     hunyuan_button.click(generate_3d_model, inputs=[segment_3d_input, rembg_Hunyuan], outputs=trellis_output)