Update app.py

app.py

@@ -4,22 +4,48 @@ import gradio as gr
 import os
 from process import load_seg_model, get_palette, generate_mask
 
-device = 'cpu'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def read_content(file_path: str) -> str:
-    with open(file_path, 'r', encoding='utf-8') as f:
-        return f.read()
+    """Read file content with error handling"""
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            return f.read()
+    except FileNotFoundError:
+        print(f"Warning: File {file_path} not found")
+        return ""
+    except Exception as e:
+        print(f"Error reading file {file_path}: {str(e)}")
+        return ""
 
 def initialize_and_load_models():
-    checkpoint_path = 'model/cloth_segm.pth'
-    return load_seg_model(checkpoint_path, device=device)
+    """Initialize and load models with error handling"""
+    try:
+        checkpoint_path = 'model/cloth_segm.pth'
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"Model checkpoint not found at {checkpoint_path}")
+        return load_seg_model(checkpoint_path, device=device)
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        return None
 
 net = initialize_and_load_models()
+if net is None:
+    raise RuntimeError("Failed to load model - check logs for details")
+
 palette = get_palette(4)
 
 def run(img):
-    cloth_seg = generate_mask(img, net=net, palette=palette, device=device)
-    return cloth_seg
+    """Process image with error handling"""
+    if img is None:
+        raise gr.Error("No image uploaded")
+    try:
+        cloth_seg = generate_mask(img, net=net, palette=palette, device=device)
+        if cloth_seg is None:
+            raise gr.Error("Failed to generate mask")
+        return cloth_seg
+    except Exception as e:
+        raise gr.Error(f"Error processing image: {str(e)}")
 
 # CSS styling
 css = '''
@@ -36,8 +62,10 @@ css = '''
 
 # Collect example images
 image_dir = 'input'
-image_list = [os.path.join(image_dir, file) for file in os.listdir(image_dir)]
-image_list.sort()
+image_list = []
+if os.path.exists(image_dir):
+    image_list = [os.path.join(image_dir, file) for file in os.listdir(image_dir) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    image_list.sort()
 examples = [[img] for img in image_list]
 
 with gr.Blocks(css=css) as demo:
@@ -57,7 +85,7 @@ with gr.Blocks(css=css) as demo:
             label="Examples - Input Images",
             examples_per_page=12
         )
-        btn = gr.Button("Run!")
+        btn = gr.Button("Run!", variant="primary")
 
     btn.click(fn=run, inputs=[image], outputs=[image_out])
 
@@ -74,5 +102,5 @@ with gr.Blocks(css=css) as demo:
         """
     )
 
-# Ensure the app works in Hugging Face by sharing a public link
-demo.launch(share=True)
+# For Hugging Face Spaces, use launch() without share=True
+demo.launch()