Update app.py

app.py

@@ -11,20 +11,23 @@ import tempfile
 import io
 import base64
 import requests
+from collections import OrderedDict
+import uuid
+
 
 # Инициализация моделей
 from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
 # oneFormer segmentation
 oneFormer_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny")
 oneFormer_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny").to(device)
-
 # classification = pipeline("image-classification", model="google/vit-base-patch16-224")
 # upscaling_client = InferenceClient(model="stabilityai/stable-diffusion-x4-upscaler")
 # inpainting_client = InferenceClient(model="stabilityai/stable-diffusion-inpainting")
+
+
 # Функции для обработки изображений
+
 def segment_image(image):
     inputs = oneFormer_processor(image, task_inputs=["panoptic"], return_tensors="pt")
 
@@ -134,7 +137,7 @@ def remove_black_make_transparent(imageEditor):
     transparent_image = Image.fromarray(image_np)
     return transparent_image
 
-def rembg(imageEditor, request: gr.Request):
+def rembg(imageEditor, rembg_model, request: gr.Request):
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
         imageEditor['composite'].save(temp_file.name)
         temp_file_path = temp_file.name
@@ -142,7 +145,7 @@ def rembg(imageEditor, request: gr.Request):
         result = client.predict(
             file=handle_file(temp_file_path),
             mask="Default",
-            model="birefnet-general-lite",
+            model=rembg_model,
             x=0,
             y=0,
             api_name="/inference"
@@ -181,6 +184,64 @@ def upscale_image(image_pil, version="v1.4", rescaling_factor=None):
     return upscaled_image
 
 
+# def inpainting(source_img, request: gr.Request):
+#     input_image = source_img["background"].convert("RGB")
+#     mask_image = source_img["layers"][0].convert("RGB")
+
+#     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=''
+):
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+        # hr_image.save(temp_file.name)
+        imageMask["background"].save(temp_file.name)
+        temp_file_path = temp_file.name #картинка
+        print("background", temp_file_path)
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file2:
+            imageMask["layers"][0].save(temp_file2.name)
+            temp_file_path2 = temp_file2.name # маска
+            print("маска", temp_file_path2)
+
+
+            client = Client("ameerazam08/FLUX.1-dev-Inpainting-Model-Beta-GPU")
+            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,
+                api_name="/process"
+            )
+            print(result)
+            return [(result, 'label')]
+
+
+            # client = Client("https://pair-hd-painter.hf.space/--replicas/0l7ng/")
+            # result = client.predict(
+            #     model_name,	# str (Option from: [('Dreamshaper Inpainting V8', 'Dreamshaper Inpainting V8'), ('Stable-Inpainting 2.0', 'Stable-Inpainting 2.0'), ('Stable-Inpainting 1.5', 'Stable-Inpainting 1.5')]) in 'Please select a model!' Dropdown component
+            #     use_rasg,	# bool  in 'Use RASG' Checkbox component
+            #     use_painta,	# bool  in 'Use PAIntA' Checkbox component
+            #     prompt,	# str  in 'Inpainting Prompt' Textbox component
+            #     handle_file(temp_file_path2),	# str (filepath on your computer (or URL) of image) in 'Input Image' Image component МАСКА
+            #     handle_file(temp_file_path),	# str (filepath on your computer (or URL) of image) in 'parameter_15' Image component КАРТИНКА
+            #     seed,	# int | float  in 'Seed' Number component
+            #     eta,	# int | float (numeric value between 0 and 1) in 'eta' Slider component
+            #     negative_prompt,	# str  in 'Negative prompt' Textbox component
+            #     positive_prompt,	# str  in 'Positive prompt' Textbox component
+            #     ddim_steps,	# int | float (numeric value between 10 and 100) in 'Number of diffusion steps' Slider component
+            #     guidance_scale,	# int | float (numeric value between 0 and 30) in 'Guidance Scale' Slider component
+            #     batch_size,	# int | float  in 'Batch size' Number component
+            #     session_id,	# str  in 'parameter_44' Textbox component
+            #     api_name="/inpaint"
+            # )
+            # print(result)
+
+
 #3d models
 def hunyuan_client(request: gr.Request):
     try:
@@ -248,6 +309,16 @@ def generate_3d_model2(image_pil, request: gr.Request):
         print(result)
         return result[0]
 
+### 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"
 
 ########## GRADIO ##########
 
@@ -276,13 +347,114 @@ with gr.Blocks() as demo:
                 with gr.Row(equal_height=True):
                     upscale_slider = gr.Slider(minimum=1, maximum=5, value=2, step=0.1, label="во сколько раз")
                     upscale_button = gr.Button("Upscale")
-                rembg_button = gr.Button("Rembg")
+                with gr.Row(equal_height=True):
+                    rembg_model_selector = gr.Dropdown(
+                        [
+                            "u2net",
+                            "u2netp",
+                            "u2net_human_seg",
+                            "u2net_cloth_seg",
+                            "silueta",
+                            "isnet-general-use",
+                            "isnet-anime",
+                            "birefnet-general",
+                            "birefnet-general-lite",
+                            "birefnet-portrait",
+                            "birefnet-dis",
+                            "birefnet-hrsod",
+                            "birefnet-cod",
+                            "birefnet-massive"
+                        ],
+                        value="birefnet-general-lite",
+                        label="Rembg model"
+                    )
+                    rembg_button = gr.Button("Rembg")
                 remove_background_button = gr.Button("Убрать черный задний фон")
                 with gr.Row(equal_height=True):
                     add_transparent_border_slider = gr.Slider(minimum=10, maximum=500, value=200, step=10, label="в пикселях")
                     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")
+                imageMask = gr.ImageEditor(
+                    label='Image',
+                    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)
+      
+                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"
+        )  
+
     with gr.Tab("Создание 3D"):
         with gr.Row(equal_height=True):
             with gr.Column(scale=5):
@@ -301,11 +473,17 @@ with gr.Blocks() as demo:
     #tab2
     crop_button.click(autocrop_image, inputs=segment_input, outputs=segment_input)
     upscale_button.click(upscale, inputs=[segment_input, upscale_slider], outputs=segment_input)
-    rembg_button.click(rembg, inputs=segment_input, outputs=segment_input)
+    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_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)
+
     #3d buttons
     hunyuan_button.click(generate_3d_model, inputs=[segment_3d_input, rembg_Hunyuan], outputs=trellis_output)
     hunyuan_button_texture.click(generate_3d_model_texture, inputs=[segment_3d_input, rembg_Hunyuan], outputs=trellis_output)