Update app.py

app.py

@@ -45,7 +45,32 @@ except Exception as e:
 
 # Model setup
 sd15_name = 'stablediffusionapi/realistic-vision-v51'
+
+# Better CUDA detection and debugging
+print("=== GPU Detection Debug ===")
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA version: {torch.version.cuda}")
+    print(f"GPU count: {torch.cuda.device_count()}")
+    print(f"Current GPU: {torch.cuda.current_device()}")
+    print(f"GPU name: {torch.cuda.get_device_name()}")
+else:
+    print("CUDA not available - checking reasons...")
+    try:
+        import subprocess
+        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
+        if result.returncode == 0:
+            print("nvidia-smi works, GPU hardware detected")
+            print("Issue might be with PyTorch CUDA installation")
+        else:
+            print("nvidia-smi failed, no GPU hardware detected")
+    except:
+        print("nvidia-smi command not found")
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Selected device: {device}")
+print("=== End GPU Debug ===")
 
 print(f"Using device: {device}")
 print("Loading models...")
@@ -354,11 +379,17 @@ def process_relight(input_fg, input_bg, prompt, image_width, image_height, num_s
     try:
         # Input validation
         if input_fg is None:
-            raise ValueError("Please upload a foreground image")
+            error_msg = "❌ Please upload a foreground image"
+            print(error_msg)
+            raise gr.Error(error_msg)
         if input_bg is None and bg_source == "Use Background Image":
-            raise ValueError("Please upload a background image or choose a lighting direction")
+            error_msg = "❌ Please upload a background image or choose a lighting direction"
+            print(error_msg)
+            raise gr.Error(error_msg)
         if not prompt.strip():
-            raise ValueError("Please enter a prompt")
+            error_msg = "❌ Please enter a prompt"
+            print(error_msg)
+            raise gr.Error(error_msg)
         
         print(f"Processing with device: {device}")
         print(f"Input shapes - FG: {input_fg.shape}, BG: {input_bg.shape if input_bg is not None else 'None'}")
@@ -392,25 +423,46 @@ def process_relight(input_fg, input_bg, prompt, image_width, image_height, num_s
                 print("Reduced num_samples to 1 for CPU processing")
         
         print("Running background removal...")
-        input_fg, matting = run_rmbg(input_fg)
+        try:
+            input_fg, matting = run_rmbg(input_fg)
+            print("Background removal completed successfully")
+        except Exception as e:
+            print(f"Background removal failed: {e}")
+            # Continue without background removal
+            matting = np.ones((input_fg.shape[0], input_fg.shape[1]), dtype=np.float32)
         
         print("Starting main processing...")
-        results, extra_images = process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
+        try:
+            results, extra_images = process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
+            print("Main processing completed successfully")
+        except Exception as e:
+            error_msg = f"❌ Processing failed: {str(e)}"
+            print(error_msg)
+            import traceback
+            traceback.print_exc()
+            raise gr.Error(error_msg)
         
         print("Converting results...")
-        results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
+        try:
+            results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
+            print("Results converted successfully")
+        except Exception as e:
+            error_msg = f"❌ Result conversion failed: {str(e)}"
+            print(error_msg)
+            raise gr.Error(error_msg)
         
         print("Processing completed successfully!")
         return results + extra_images
         
+    except gr.Error:
+        # Re-raise Gradio errors to show them in the UI
+        raise
     except Exception as e:
-        print(f"Error in process_relight: {str(e)}")
+        error_msg = f"❌ Unexpected error: {str(e)}"
+        print(error_msg)
         import traceback
         traceback.print_exc()
-        # Return error image
-        error_img = np.zeros((512, 512, 3), dtype=np.uint8)
-        error_img[:, :] = [255, 0, 0]  # Red error image
-        return [error_img]
+        raise gr.Error(error_msg)
 
 # Quick prompts for easy testing
 quick_prompts = [