Add main space app.py update

app.py

@@ -1,23 +1,33 @@
+import torch
 import numpy as np
 import cv2
-import torch
+from PIL import Image
 import albumentations as albu
 from pylab import imshow
 import matplotlib.pyplot as plt
 from diffusers import StableDiffusionInpaintPipeline
-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")
-model.eval()
+try:
+    model = create_model("Unet_2020-10-30")
+    model.eval()
+except Exception as e:
+    raise RuntimeError(f"Error loading segmentation model: {e}")
+
+# Load Inpainting Model
+try:
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.float16
+    )
+    pipe.to("cuda")  
+except Exception as e:
+    raise RuntimeError(f"Error loading inpainting model: {e}")
 
-# 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)
@@ -36,7 +46,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 = Image.open(mask_path).convert("L") 
     mask_image = mask_image.resize(image.size)
     output_image = pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
     return output_image
@@ -47,25 +57,25 @@ def resize_and_upscale(image, new_width, new_height):
 
 
 def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armor"):
-    pil_image = Image.fromarray(image.astype('uint8'))  # Gradio provides image as np.array
+    try:
+        pil_image = Image.fromarray(image.astype('uint8'))  
+        temp_image_path = "temp_image.jpg"
+        pil_image.save(temp_image_path)
 
-    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)
 
-    x, image, pads = load_and_preprocess_image(temp_image_path)
-    mask = segment_cloth(x, model, pads)
+        mask_path = "temp_mask.jpg"
+        plt.imsave(mask_path, mask, cmap='gray')
 
-    mask_path = "temp_mask.jpg"
-    plt.imsave(mask_path, mask, cmap='gray')
-
-    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
+        output_image = perform_inpainting(temp_image_path, mask_path, prompt)
+        output_image = resize_and_upscale(output_image, 1280, 720) 
+        return output_image
+    except Exception as e:
+        raise gr.Error(f"Error processing image: {e}")
 
 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")

requirements.txt

@@ -1,10 +1,11 @@
 torch
+accelerate
 albumentations
 matplotlib
 diffusers
 transformers
-iglovikov_helper_functions  # (Replace with the correct installation method if not on PyPI)
-cloths_segmentation         # (Replace with the correct installation method if not on PyPI)
+iglovikov_helper_functions  
+cloths_segmentation  
 opencv-python-headless
 gradio
 numpy==1.23.5  # Example version; adjust if necessary