Update app.py

app.py

@@ -11,31 +11,36 @@ import huggingface_hub
 huggingface_hub.login(os.getenv('HF_TOKEN_FLUX'))
 
 check_min_version("0.30.2")
-
-
-# Build pipeline
-controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
 transformer = FluxTransformer2DModel.from_pretrained(
         "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dytpe=torch.bfloat16
     )
-pipe = FluxControlNetInpaintingPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    controlnet=controlnet,
-    transformer=transformer,
-    torch_dtype=torch.bfloat16
-).to("cuda")
-pipe.transformer.to(torch.bfloat16)
-pipe.controlnet.to(torch.bfloat16)
+
+
+# Build pipeline
+
+def load_models(model_name):
+    controlnet = FluxControlNetModel.from_pretrained(f"alimama-creative/{model_name}", torch_dtype=torch.bfloat16)
+    pipe = FluxControlNetInpaintingPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        controlnet=controlnet,
+        transformer=transformer,
+        torch_dtype=torch.bfloat16
+    ).to("cuda")
+    pipe.transformer.to(torch.bfloat16)
+    pipe.controlnet.to(torch.bfloat16)
+    return pipe
 
 
 MARKDOWN = """
-# FLUX.1-dev-Inpainting-Model-BETA-GPU 🔥
+# FLUX.1-dev-Inpainting-Model-GPU 🔥
+Model by alimama-creative
 """
 
 @spaces.GPU()
 def process(input_image_editor,
             prompt,
             negative_prompt,
+            model_name,
             controlnet_conditioning_scale,
             guidance_scale,
             seed,
@@ -46,6 +51,8 @@ def process(input_image_editor,
     mask = input_image_editor['layers'][0]
     size = (768, 768)
     image_or = image.copy()
+    pipe = load_models(model_name)
+    pipe =  pipe.to("cuda")
     
     image = image.convert("RGB").resize(size)
     mask = mask.convert("RGB").resize(size)
@@ -80,6 +87,7 @@ with gr.Blocks() as demo:
 
             prompt = gr.Textbox(lines=2, placeholder="Enter prompt here...")
             negative_prompt = gr.Textbox(lines=2, placeholder="Enter negative_prompt here...")
+            model_name = gr.Dropdown(choices=["FLUX.1-dev-Controlnet-Inpainting-Beta","FLUX.1-dev-Controlnet-Inpainting-Alpha"], label="Select Model Inpainting", value="FLUX.1-dev-Controlnet-Inpainting-Beta"),
             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")
@@ -101,6 +109,7 @@ with gr.Blocks() as demo:
             input_image_editor_component,
             prompt,
             negative_prompt,
+            model_name,
             controlnet_conditioning_scale,
             guidance_scale,
             seed,