Create app.py

app.py

@@ -1,146 +1,165 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionInpaintPipeline, StableDiffusionXLInpaintPipeline
+import gc
+from diffusers import StableDiffusionInpaintPipeline
 from PIL import Image
 import numpy as np
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
 import cv2
+import traceback
 
 class RafayyVirtualTryOn:
     def __init__(self):
-        # Initialize SDXL for better quality
-        self.inpaint_model = StableDiffusionXLInpaintPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            variant="fp16"
-        )
-        if torch.cuda.is_available():
-            self.inpaint_model.to("cuda")
+        try:
+            # Clear CUDA cache
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                gc.collect()
+
+            # Use smaller model for stability
+            self.inpaint_model = StableDiffusionInpaintPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                safety_checker=None  # Disable safety checker if causing issues
+            )
+            
+            if torch.cuda.is_available():
+                self.inpaint_model.to("cuda")
+                self.inpaint_model.enable_attention_slicing()  # Reduce memory usage
             
-        # Initialize enhanced segmentation model
-        self.segmenter = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
-        self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
-        
-    def enhance_mask(self, mask):
-        """Enhance the clothing mask for better results"""
-        kernel = np.ones((5,5), np.uint8)
-        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
-        mask = cv2.GaussianBlur(mask, (5,5), 0)
-        return mask
+            # Initialize segmentation with error handling
+            try:
+                self.segmenter = SegformerForSemanticSegmentation.from_pretrained(
+                    "mattmdjaga/segformer_b2_clothes",
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+                )
+                self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+            except Exception as e:
+                print(f"Segmentation model loading error: {str(e)}")
+                raise
+                
+        except Exception as e:
+            print(f"Initialization error: {str(e)}")
+            raise
 
-    def get_clothing_mask(self, image):
-        """Extract and enhance clothing mask"""
-        inputs = self.processor(images=image, return_tensors="pt")
-        outputs = self.segmenter(**inputs)
-        logits = outputs.logits.squeeze()
-        
-        clothing_mask = (logits.argmax(0) == 5).float().numpy()
-        clothing_mask = (clothing_mask * 255).astype(np.uint8)
-        clothing_mask = self.enhance_mask(clothing_mask)
-        
-        return Image.fromarray(clothing_mask)
+    def preprocess_image(self, image):
+        """Safely preprocess input image"""
+        try:
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            
+            # Ensure image is RGB
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            
+            # Resize if too large
+            max_size = 768
+            if max(image.size) > max_size:
+                ratio = max_size / max(image.size)
+                new_size = tuple(int(dim * ratio) for dim in image.size)
+                image = image.resize(new_size, Image.LANCZOS)
+            
+            return image
+        except Exception as e:
+            raise gr.Error(f"Image preprocessing failed: {str(e)}")
 
-    def try_on(self, original_image, prompt, style_strength=0.7, progress=gr.Progress()):
-        """Enhanced virtual try-on with style control"""
+    def get_clothing_mask(self, image):
+        """Safely extract clothing mask"""
         try:
-            progress(0, desc="Initializing...")
+            # Convert to RGB if needed
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+
+            inputs = self.processor(images=image, return_tensors="pt")
             
-            # Image preprocessing
-            original_image = Image.fromarray(original_image)
-            if original_image.mode != "RGB":
-                original_image = original_image.convert("RGB")
+            # Move to GPU if available
+            if torch.cuda.is_available():
+                inputs = {k: v.to("cuda") for k, v in inputs.items()}
+                self.segmenter = self.segmenter.to("cuda")
+
+            outputs = self.segmenter(**inputs)
+            logits = outputs.logits.squeeze()
             
-            progress(0.2, desc="Analyzing clothing...")
-            mask = self.get_clothing_mask(original_image)
+            # Move back to CPU for numpy operations
+            if torch.cuda.is_available():
+                logits = logits.cpu()
+
+            clothing_mask = (logits.argmax(0) == 5).float().numpy()
+            clothing_mask = (clothing_mask * 255).astype(np.uint8)
             
-            # Enhanced prompt engineering
-            style_prompts = {
-                "quality": "ultra detailed, 8k uhd, high quality, professional photo",
-                "lighting": "perfect lighting, studio lighting, professional photography",
-                "realism": "hyperrealistic, photorealistic, highly detailed"
-            }
+            # Enhance mask
+            kernel = np.ones((5,5), np.uint8)
+            clothing_mask = cv2.dilate(clothing_mask, kernel, iterations=1)
+            clothing_mask = cv2.GaussianBlur(clothing_mask, (5,5), 0)
+            
+            return Image.fromarray(clothing_mask)
+        except Exception as e:
+            raise gr.Error(f"Mask generation failed: {str(e)}")
+
+    def try_on(self, image, prompt, style_strength=0.7, progress=gr.Progress()):
+        """Main try-on function with comprehensive error handling"""
+        try:
+            if image is None:
+                raise gr.Error("Please upload an image first")
+            if not prompt or prompt.strip() == "":
+                raise gr.Error("Please provide a clothing description")
+
+            progress(0.1, desc="Preprocessing image...")
+            original_image = self.preprocess_image(image)
             
-            full_prompt = f"""
-            A person wearing {prompt}, {style_prompts['quality']}, 
-            {style_prompts['lighting']}, {style_prompts['realism']}
-            """
+            progress(0.3, desc="Detecting clothing...")
+            mask = self.get_clothing_mask(original_image)
             
-            negative_prompt = """
-            low quality, blurry, distorted, deformed, unrealistic, 
-            bad proportions, bad lighting, oversaturated, undersaturated
-            """
+            # Clear GPU memory
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                gc.collect()
             
-            progress(0.4, desc="Generating new clothing...")
-            result = self.inpaint_model(
-                prompt=full_prompt,
-                negative_prompt=negative_prompt,
-                image=original_image,
-                mask_image=mask,
-                num_inference_steps=50,
-                guidance_scale=7.5 * style_strength
-            ).images[0]
+            progress(0.5, desc="Preparing generation...")
+            # Enhanced prompt engineering
+            full_prompt = f"A person wearing {prompt}, professional photo, detailed, realistic, high quality"
+            negative_prompt = "low quality, blurry, distorted, deformed, bad anatomy, unrealistic"
             
-            progress(1.0, desc="Finalizing...")
+            progress(0.7, desc="Generating new clothing...")
+            try:
+                result = self.inpaint_model(
+                    prompt=full_prompt,
+                    negative_prompt=negative_prompt,
+                    image=original_image,
+                    mask_image=mask,
+                    num_inference_steps=30,  # Reduced for stability
+                    guidance_scale=7.5 * style_strength
+                ).images[0]
+            except torch.cuda.OutOfMemoryError:
+                torch.cuda.empty_cache()
+                gc.collect()
+                raise gr.Error("Out of memory. Please try with a smaller image.")
+            
+            progress(1.0, desc="Done!")
             return result
             
         except Exception as e:
-            raise gr.Error(f"Processing Error: {str(e)}")
-
-# Initialize the model
-model = RafayyVirtualTryOn()
+            error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+            print(error_msg)  # For logging
+            raise gr.Error(str(e))
 
-# Custom CSS with professional styling
-custom_css = """
-.gradio-container {
-    font-family: 'Poppins', sans-serif;
-    max-width: 1200px !important;
-    margin: auto !important;
-}
-#component-0 {
-    max-width: 100% !important;
-    margin-bottom: 20px !important;
-}
-.gr-button {
-    background: linear-gradient(90deg, #2193b0, #6dd5ed) !important;
-    border: none !important;
-    color: white !important;
-    font-weight: 600 !important;
-}
-.gr-button:hover {
-    background: linear-gradient(90deg, #6dd5ed, #2193b0) !important;
-    transform: translateY(-2px);
-    box-shadow: 0 5px 15px rgba(33, 147, 176, 0.3) !important;
-    transition: all 0.3s ease;
-}
-.gr-input {
-    border: 2px solid #e0e0e0 !important;
-    border-radius: 8px !important;
-    padding: 12px !important;
-}
-.gr-input:focus {
-    border-color: #2193b0 !important;
-    box-shadow: 0 0 0 2px rgba(33, 147, 176, 0.2) !important;
-}
-.gr-panel {
-    border-radius: 12px !important;
-    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1) !important;
-}
-.footer {
-    background: linear-gradient(to right, #f8f9fa, #e9ecef);
-    padding: 20px;
-    border-radius: 10px;
-    margin-top: 30px;
-}
-"""
+# Initialize model with error handling
+try:
+    model = RafayyVirtualTryOn()
+except Exception as e:
+    print(f"Model initialization failed: {str(e)}")
+    raise
 
-# Create enhanced Gradio interface
+# Create Gradio interface with error handling
 demo = gr.Interface(
     fn=model.try_on,
     inputs=[
         gr.Image(label="📸 Upload Your Photo", type="numpy"),
         gr.Textbox(
             label="🎨 Describe New Clothing",
-            placeholder="e.g., 'elegant black suit with silk lapels', 'designer red dress with gold accents'",
+            placeholder="e.g., 'elegant black suit', 'red dress'",
             lines=2
         ),
         gr.Slider(
@@ -151,59 +170,32 @@ demo = gr.Interface(
             step=0.1
         )
     ],
-    outputs=gr.Image(label="✨ Your New Look", type="pil"),
-    title="🌟 Rafayy's Professional Virtual Try-On Studio 🌟",
+    outputs=gr.Image(label="✨ Result", type="pil"),
+    title="🌟 Rafayy's Virtual Try-On Studio 🌟",
     description="""
-    <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
-        <h3 style="color: #2193b0; font-size: 24px; margin-bottom: 10px;">Transform Your Style with AI</h3>
-        <p style="color: #666; font-size: 16px;">Experience the future of fashion with our advanced virtual try-on technology.</p>
-        <div style="margin-top: 20px; padding: 15px; background: rgba(33, 147, 176, 0.1); border-radius: 10px;">
-            <p style="font-weight: bold; color: #2193b0;">Premium Features:</p>
-            <p>✓ High-Resolution Output</p>
-            <p>✓ Advanced Clothing Detection</p>
-            <p>✓ Professional Style Enhancement</p>
-        </div>
+    <div style="text-align: center;">
+        <h3>Transform Your Style with AI</h3>
+        <p>Upload a photo and describe the new clothing you want to try on!</p>
     </div>
     """,
     examples=[
-        ["example1.jpg", "luxury black suit with silk details", 0.8],
-        ["example2.jpg", "designer white dress with lace accents", 0.7],
-        ["example3.jpg", "professional navy blazer with gold buttons", 0.9],
-        ["example4.jpg", "casual denim jacket with vintage wash", 0.6]
+        ["example1.jpg", "black suit", 0.7],
+        ["example2.jpg", "white dress", 0.7]
     ],
-    css=custom_css
+    allow_flagging="never",
+    cache_examples=True
 )
 
-# Enhanced footer with professional information
-demo.footer = """
-<div class="footer">
-    <div style="text-align: center;">
-        <h4 style="color: #2193b0; margin-bottom: 15px;">🎯 Professional Tips for Best Results</h4>
-        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 20px;">
-            <div>
-                <p style="font-weight: bold;">📸 Photo Guidelines</p>
-                <p>• Use well-lit, front-facing photos</p>
-                <p>• Ensure clear visibility of clothing</p>
-                <p>• Avoid complex backgrounds</p>
-            </div>
-            <div>
-                <p style="font-weight: bold;">✍️ Description Tips</p>
-                <p>• Be specific about style details</p>
-                <p>• Include color and material</p>
-                <p>• Mention design elements</p>
-            </div>
-            <div>
-                <p style="font-weight: bold;">⚙️ Settings</p>
-                <p>• Adjust style strength as needed</p>
-                <p>• Higher values for bold changes</p>
-                <p>• Lower values for subtle effects</p>
-            </div>
-        </div>
-        <p style="margin-top: 20px; color: #666;">© 2024 Rafayy AI Studio | Professional Virtual Try-On Service</p>
-    </div>
-</div>
-"""
-
-# Launch the app
+# Launch with error handling
 if __name__ == "__main__":
-    demo.launch()
+    try:
+        demo.launch(
+            share=False,
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True,
+            enable_queue=True
+        )
+    except Exception as e:
+        print(f"Launch failed: {str(e)}")
+        raise