GPU error

app.py

@@ -5,162 +5,243 @@ from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentati
 import numpy as np
 from diffusers import StableDiffusionInpaintPipeline
 import warnings
+import os
 warnings.filterwarnings("ignore")
 
+# Force CPU usage to avoid GPU issues on Hugging Face Spaces
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+torch.set_default_dtype(torch.float32)
+
 # Global variables for models
 processor = None
 model = None
 pipe = None
 
+def get_device():
+    """Safely determine the best available device"""
+    try:
+        # Force CPU for stability on HF Spaces
+        return "cpu"
+    except:
+        return "cpu"
+
 def load_models():
-    """Load models with proper error handling"""
+    """Load models with CPU-only configuration"""
     global processor, model, pipe
     
     try:
         print("Loading segmentation model...")
         processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
         model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+        
+        # Ensure segmentation model is on CPU
+        model = model.to("cpu")
+        model.eval()
         print("Segmentation model loaded successfully!")
         
         print("Loading Stable Diffusion inpainting model...")
-        # Use CPU for Hugging Face Spaces (most don't have GPU access)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Load with explicit CPU configuration
         pipe = StableDiffusionInpaintPipeline.from_pretrained(
             "stabilityai/stable-diffusion-2-inpainting",
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            safety_checker=None,  # Disable safety checker to save memory
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            safety_checker=None,
             requires_safety_checker=False,
             use_safetensors=True
         )
-        pipe = pipe.to(device)
         
-        # Enable memory efficient attention if available
+        # Explicitly move all components to CPU
+        pipe = pipe.to("cpu")
+        
+        # Enable memory efficiency
         if hasattr(pipe, 'enable_attention_slicing'):
             pipe.enable_attention_slicing()
-        if hasattr(pipe, 'enable_model_cpu_offload'):
-            pipe.enable_model_cpu_offload()
+        
+        # Set to eval mode
+        pipe.unet.eval()
+        pipe.vae.eval()
+        if hasattr(pipe, 'text_encoder'):
+            pipe.text_encoder.eval()
             
-        print("Models loaded successfully!")
+        print("Stable Diffusion model loaded successfully on CPU!")
         return True
         
     except Exception as e:
         print(f"Error loading models: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return False
 
 def segment_clothes(human_image):
-    """Segment clothing from human image"""
+    """Segment clothing from human image with CPU-only operations"""
     try:
-        # Resize image if too large to save memory
+        # Resize image if too large
+        original_size = human_image.size
         if human_image.size[0] > 512 or human_image.size[1] > 512:
             human_image = human_image.resize((512, 512), Image.Resampling.LANCZOS)
         
         # Process human image for segmentation
         inputs = processor(images=human_image, return_tensors="pt")
         
+        # Ensure inputs are on CPU
+        for key in inputs:
+            if torch.is_tensor(inputs[key]):
+                inputs[key] = inputs[key].to("cpu")
+        
         with torch.no_grad():
             outputs = model(**inputs)
             
         logits = outputs.logits.cpu()
         upsampled_logits = torch.nn.functional.interpolate(
-            logits, size=human_image.size[::-1], mode="bilinear", align_corners=False
+            logits, 
+            size=human_image.size[::-1], 
+            mode="bilinear", 
+            align_corners=False
         )
         pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
         
-        # Create mask for clothes (labels 4, 5, 6, 7 typically represent different clothing items)
-        clothes_mask = np.isin(pred_seg, [4, 5, 6, 7]).astype(np.uint8) * 255
+        # Create mask for clothes
+        clothes_labels = [4, 5, 6, 7, 8, 9, 10]
+        clothes_mask = np.isin(pred_seg, clothes_labels).astype(np.uint8) * 255
+        
+        # If no clothes detected, create a default mask
+        if np.sum(clothes_mask) < 100:
+            print("Creating default upper body mask")
+            mask = np.zeros_like(pred_seg, dtype=np.uint8)
+            h, w = mask.shape
+            # Upper body region
+            mask[h//4:3*h//4, w//3:2*w//3] = 255
+            clothes_mask = mask
         
-        return Image.fromarray(clothes_mask)
+        # Resize back to original size
+        mask_image = Image.fromarray(clothes_mask)
+        if original_size != mask_image.size:
+            mask_image = mask_image.resize(original_size, Image.Resampling.LANCZOS)
+        
+        return mask_image
         
     except Exception as e:
         print(f"Error in segmentation: {str(e)}")
-        # Return a default mask if segmentation fails
-        return Image.new('L', human_image.size, 255)
+        # Return a default center mask
+        h, w = human_image.size[::-1]
+        mask = np.zeros((h, w), dtype=np.uint8)
+        mask[h//4:3*h//4, w//3:2*w//3] = 255
+        return Image.fromarray(mask)
 
 def try_on_cloth(human_image, cloth_image, progress=gr.Progress()):
-    """Main function for virtual try-on"""
+    """Main function for virtual try-on with CPU-safe operations"""
     if human_image is None or cloth_image is None:
         return None, "Please upload both human and cloth images."
     
     if processor is None or model is None or pipe is None:
-        return None, "Models not loaded. Please wait for initialization to complete."
+        return None, "Models not loaded. Please refresh the page and try again."
     
     try:
         progress(0.1, desc="Processing images...")
         
-        # Ensure images are PIL Images and resize for memory efficiency
+        # Ensure images are PIL Images
         if isinstance(human_image, np.ndarray):
             human_image = Image.fromarray(human_image)
         if isinstance(cloth_image, np.ndarray):
             cloth_image = Image.fromarray(cloth_image)
-            
-        # Resize images to reasonable size for processing
+        
+        # Convert to RGB
+        if human_image.mode != 'RGB':
+            human_image = human_image.convert('RGB')
+        if cloth_image.mode != 'RGB':
+            cloth_image = cloth_image.convert('RGB')
+        
+        # Resize for processing
         target_size = (512, 512)
         human_image = human_image.resize(target_size, Image.Resampling.LANCZOS)
         cloth_image = cloth_image.resize(target_size, Image.Resampling.LANCZOS)
         
         progress(0.3, desc="Generating clothing mask...")
         
-        # Generate clothes mask
+        # Generate mask
         mask = segment_clothes(human_image)
         
-        progress(0.6, desc="Generating try-on result...")
+        progress(0.6, desc="Generating try-on result (this may take a few minutes on CPU)...")
         
-        # Use Stable Diffusion Inpainting
-        prompt = "A person wearing the provided clothing, realistic, high quality, detailed"
-        negative_prompt = "blurry, low quality, distorted, deformed"
+        # Prepare for inpainting
+        prompt = "a person wearing the clothing, realistic, high quality, natural lighting"
+        negative_prompt = "blurry, low quality, distorted, deformed, extra limbs"
         
-        # Generate result with lower parameters for faster processing
-        result = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=human_image,
-            mask_image=mask,
-            num_inference_steps=20,  # Reduced for faster processing
-            strength=0.8,
-            guidance_scale=7.5,
-            generator=torch.Generator().manual_seed(42)  # For reproducible results
-        ).images[0]
+        # Create CPU generator
+        generator = torch.Generator(device='cpu').manual_seed(42)
         
-        progress(1.0, desc="Complete!")
+        # Generate with CPU-optimized settings
+        with torch.no_grad():
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=human_image,
+                mask_image=mask,
+                num_inference_steps=15,  # Reduced for CPU
+                strength=0.75,
+                guidance_scale=7.0,
+                generator=generator
+            ).images[0]
         
-        return result, "Try-on completed successfully!"
+        progress(1.0, desc="Complete!")
+        return result, "Try-on completed successfully! (Processed on CPU)"
         
     except Exception as e:
         error_msg = f"Error during try-on: {str(e)}"
         print(error_msg)
-        return None, error_msg
+        import traceback
+        traceback.print_exc()
+        
+        # Attempt simple fallback
+        try:
+            progress(0.8, desc="Attempting simple blend fallback...")
+            mask_array = np.array(mask) / 255.0
+            cloth_resized = cloth_image.resize(human_image.size)
+            
+            human_array = np.array(human_image).astype(np.float32)
+            cloth_array = np.array(cloth_resized).astype(np.float32)
+            
+            mask_3d = np.stack([mask_array] * 3, axis=2)
+            result_array = human_array * (1 - mask_3d) + cloth_array * mask_3d
+            result = Image.fromarray(result_array.astype(np.uint8))
+            
+            return result, "Used simple blending due to processing error."
+        except:
+            return None, error_msg
 
-# Initialize models when the app starts
-print("Initializing models...")
+# Initialize models
+print("Initializing models for CPU processing...")
 models_loaded = load_models()
 
-# Create Gradio interface
-with gr.Blocks(title="Virtual Cloth Try-On AI", theme=gr.themes.Soft()) as interface:
+# Gradio interface
+with gr.Blocks(title="Virtual Cloth Try-On AI", theme=gr.themes.Default()) as interface:
     gr.Markdown("""
-    # 🧥 Virtual Cloth Try-On AI
+    # 🧥 Virtual Cloth Try-On AI (CPU Version)
     
     Upload a photo of a person and a clothing item to see how the outfit would look!
     
-    **Instructions:**
-    1. Upload a clear photo of a person (front-facing works best) in the first box
-    2. Upload an image of the clothing item you want to try on in the second box
-    3. Click "Generate Try-On" and wait for the result
+    **⚠️ Note: This app runs on CPU, so processing will take 2-5 minutes per image.**
     
-    **Note:** Processing may take 1-2 minutes depending on server load.
+    **Instructions:**
+    1. Upload a clear photo of a person (front-facing works best)
+    2. Upload an image of the clothing item you want to try on
+    3. Click "Generate Try-On" and be patient - CPU processing is slow but works!
     """)
     
     if not models_loaded:
-        gr.Markdown("⚠️ **Models failed to load. Please refresh the page.**")
+        gr.Markdown("❌ **Models failed to load. Please refresh the page.**")
+    else:
+        gr.Markdown("✅ **Models loaded successfully! Ready for try-on.**")
     
     with gr.Row():
         with gr.Column():
             human_input = gr.Image(
                 type="pil",
-                label="👤 Human Image"
+                label="👤 Human Photo"
             )
             cloth_input = gr.Image(
                 type="pil", 
-                label="👕 Clothing Image"
+                label="👕 Clothing Item"
             )
             
         with gr.Column():
@@ -170,11 +251,12 @@ with gr.Blocks(title="Virtual Cloth Try-On AI", theme=gr.themes.Soft()) as inter
             )
             status_output = gr.Textbox(
                 label="Status",
-                interactive=False
+                interactive=False,
+                placeholder="Upload images and click 'Generate Try-On'"
             )
     
     generate_btn = gr.Button(
-        "🎨 Generate Try-On",
+        "🎨 Generate Try-On (Takes 2-5 minutes)",
         variant="primary",
         size="lg"
     )
@@ -189,10 +271,13 @@ with gr.Blocks(title="Virtual Cloth Try-On AI", theme=gr.themes.Soft()) as inter
     gr.Markdown("""
     ---
     **Tips for better results:**
-    - Use high-resolution, well-lit images
-    - Ensure the person is facing forward
-    - Use clothing images with clear, visible details
-    - Avoid complex backgrounds when possible
+    - Use clear, high-resolution images with good lighting
+    - Person should be facing forward with visible torso
+    - Clothing items should be clearly visible and unfolded
+    - Simple backgrounds work better than busy ones
+    - Be patient - CPU processing takes time but produces good results!
+    
+    **Expected processing time: 2-5 minutes per try-on**
     """)
 
 if __name__ == "__main__":