Update app.py

app.py

@@ -8,31 +8,25 @@ import os
 import shutil
 from gradio_client import Client, handle_file
 
-# Function to load Stable Diffusion Inpainting model
-def load_inpainting_model(model_name):
-    model_path = model_name
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load the model once globally to avoid repeated loading
+def load_inpainting_model():
+    model_path = "uberRealisticPornMerge_v13Inpainting.safetensors"
+    device = "cpu"  # Explicitly use CPU
     pipe = StableDiffusionInpaintPipeline.from_single_file(
         model_path,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        torch_dtype=torch.float32,  # Use float32 for CPU
         safety_checker=None
     ).to(device)
     return pipe
 
-# Function to resize image
+# Preload the model once
+inpaint_pipeline = load_inpainting_model()
+
+# Function to resize image (simpler interpolation method for speed)
 def resize_to_match(input_image, output_image):
-    return output_image.resize(input_image.size, Image.LANCZOS)
-
-# Function to upscale image
-def upscale_image(image):
-    scale_factor = 2 if torch.cuda.is_available() else 1
-    image_np = np.array(image)
-    if scale_factor == 1:
-        return image
-    upscaled_image_np = cv2.resize(image_np, (0, 0), fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_CUBIC)
-    return Image.fromarray(upscaled_image_np)
-
-# Step 1: Generate the mask using Florence SAM Masking API (Replicate)
+    return output_image.resize(input_image.size, Image.BILINEAR)  # Use BILINEAR for faster resizing
+
+# Function to generate the mask using Florence SAM Masking API (Replicate)
 def generate_mask(image_path, text_prompt="clothing"):
     client_sam = Client("SkalskiP/florence-sam-masking")
     mask_result = client_sam.predict(
@@ -42,49 +36,41 @@ def generate_mask(image_path, text_prompt="clothing"):
     )
     return mask_result  # This is the local path to the generated mask
 
-# Step 2: Save the generated mask to the same directory
+# Save the generated mask
 def save_mask(mask_local_path, save_path="generated_mask.png"):
     try:
         shutil.copy(mask_local_path, save_path)
-        print(f"Generated mask saved to {save_path}")
     except Exception as e:
         print(f"Failed to save the mask: {e}")
 
 # Function to perform inpainting
-def inpaint_image(input_image, mask_image, inpaint_pipeline):
+def inpaint_image(input_image, mask_image):
     prompt = "undress, naked"
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     result = inpaint_pipeline(prompt=prompt, image=input_image, mask_image=mask_image)
     inpainted_image = result.images[0]
     inpainted_image = resize_to_match(input_image, inpainted_image)
-    inpainted_image = upscale_image(inpainted_image)
     return inpainted_image
 
-# Function to process input image and mask with the selected model
-def process_image(input_image, model_choice):
+# Function to process input image and mask
+def process_image(input_image):
     # Save the input image temporarily to process with Replicate
     input_image_path = "temp_input_image.png"
     input_image.save(input_image_path)
 
-    # Step 1: Generate the mask using the Florence SAM Masking API
+    # Generate the mask using Florence SAM API
     mask_local_path = generate_mask(image_path=input_image_path)
 
-    # Step 2: Save the generated mask
+    # Save the generated mask
     mask_image_path = "generated_mask.png"
     save_mask(mask_local_path, save_path=mask_image_path)
 
-    # Step 3: Load the selected inpainting model
-    inpaint_pipeline = load_inpainting_model(model_choice)
-
-    # Step 4: Open the mask image and perform inpainting
+    # Open the mask image and perform inpainting
     mask_image = Image.open(mask_image_path)
-    result_image = inpaint_image(input_image, mask_image, inpaint_pipeline)
+    result_image = inpaint_image(input_image, mask_image)
 
-    # Step 5: Clean up temporary files
-    if os.path.exists(input_image_path):
-        os.remove(input_image_path)
-    if os.path.exists(mask_image_path):
-        os.remove(mask_image_path)
+    # Clean up temporary files
+    os.remove(input_image_path)
+    os.remove(mask_image_path)
 
     return result_image
 
@@ -94,22 +80,12 @@ with gr.Blocks() as demo:
         input_image = gr.Image(label="Upload Input Image", type="pil")
         output_image = gr.Image(type="pil", label="Output Image")
 
-    # Dropdown for selecting the model checkpoint
-    model_choice = gr.Dropdown(
-        choices=[
-            "uberRealisticPornMerge_v13Inpainting.safetensors", 
-            "realvisxlInpainting_v50Inpainting.safetensors"
-        ],
-        label="Choose Inpainting Model",
-        value="uberRealisticPornMerge_v13Inpainting.safetensors"  # Default model
-    )
-
     # Button to trigger the process
     with gr.Row():
         btn = gr.Button("Run Inpainting")
 
     # Function to run when button is clicked
-    btn.click(fn=process_image, inputs=[input_image, model_choice], outputs=output_image)
+    btn.click(fn=process_image, inputs=[input_image], outputs=output_image)
 
-# Launch the Gradio app with share=True to allow public access
+# Launch the Gradio app
 demo.launch(share=True)