Add main space file

app.py

@@ -18,6 +18,7 @@ try:
 except Exception as e:
     raise RuntimeError(f"Error loading segmentation model: {e}")
 
+
 # Load Inpainting Model (Without CUDA)
 try:
     pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
@@ -28,9 +29,7 @@ except Exception as e:
 def load_and_preprocess_image(image_path):
     image = load_rgb(image_path)
     padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
-    transform = albu.Compose([albu.Normalize(p=1)], p=1)
-    x = transform(image=padded_image)["image"]
-    x = torch.unsqueeze(tensor_from_rgb_image(x), 0)
+    x = torch.unsqueeze(tensor_from_rgb_image(image), 0)  # Use original image, not padded
     return x, image, pads
 
 def segment_cloth(image_tensor, model, pads):
@@ -65,6 +64,7 @@ def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armo
 
         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
     except Exception as e:
         raise gr.Error(f"Error processing image: {e}")
@@ -81,4 +81,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)