Add main space app.py file

app.py

@@ -9,6 +9,7 @@ from PIL import Image
 from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
 from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
 from cloths_segmentation.pre_trained_models import create_model
+import gradio as gr
 
 # Load Cloth Segmentation Model (Ensure this is available)
 model = create_model("Unet_2020-10-30")
@@ -16,6 +17,7 @@ model.eval()
 
 # Load Inpainting Model (Ensure this is available)
 pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
+pipe.to("cuda")  # If you have a GPU, this will run the model on it
 
 def load_and_preprocess_image(image_path):
     image = load_rgb(image_path)
@@ -35,8 +37,7 @@ def segment_cloth(image_tensor, model, pads):
 def perform_inpainting(image_path, mask_path, prompt):
     image = Image.open(image_path)
     mask_image = Image.open(mask_path).convert("L") # Convert to single-channel grayscale
-    mask_image = mask_image.resize(image.size)  # Resize mask to match image
-
+    mask_image = mask_image.resize(image.size)
     output_image = pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
     return output_image
 
@@ -44,24 +45,27 @@ def resize_and_upscale(image, new_width, new_height):
     resized_img = cv2.resize(np.array(image), (new_width, new_height), interpolation=cv2.INTER_CUBIC)
     return Image.fromarray(resized_img)
 
-import gradio as gr
 
 def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armor"):
-    x, image, pads = load_and_preprocess_image(image.name)  # Gradio provides image.name for the path
+    pil_image = Image.fromarray(image.astype('uint8'))  # Gradio provides image as np.array
+
+    temp_image_path = "temp_image.jpg"
+    pil_image.save(temp_image_path)
+
+    x, image, pads = load_and_preprocess_image(temp_image_path)
     mask = segment_cloth(x, model, pads)
 
-    # Save mask temporarily
     mask_path = "temp_mask.jpg"
     plt.imsave(mask_path, mask, cmap='gray')
 
-    output_image = perform_inpainting(image.name, mask_path, prompt)
-    output_image = resize_and_upscale(output_image, 1280, 720)  # Adjust dimensions as needed
-    return output_image
+    output_image = perform_inpainting(temp_image_path, mask_path, prompt)
+    output_image = resize_and_upscale(output_image, 1280, 720)  # You can adjust the size
 
+    return output_image
 
 with gr.Blocks() as demo:
     gr.Markdown("# Cloth Image Segmentation and Inpainting")
-    
+
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(label="Upload Image")
@@ -72,4 +76,4 @@ with gr.Blocks() as demo:
 
     run_button.click(fn=image_segmentation_and_inpainting, inputs=[image_input, prompt_input], outputs=image_output)
 
-demo.launch(share=True) 
+demo.launch(share=True)