Create app.py

app.py

@@ -7,6 +7,12 @@ import numpy as np
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
 import cv2
 import traceback
+import os
+
+# Set environment variables to prevent warnings and errors
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ['TORCH_HOME'] = '/tmp/torch'
+os.environ['HF_HOME'] = '/tmp/huggingface'
 
 class RafayyVirtualTryOn:
     def __init__(self):
@@ -17,23 +23,29 @@ class RafayyVirtualTryOn:
                 gc.collect()
 
             # Use smaller model for stability
+            model_id = "runwayml/stable-diffusion-inpainting"
             self.inpaint_model = StableDiffusionInpaintPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5",
+                model_id,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                safety_checker=None  # Disable safety checker if causing issues
+                safety_checker=None,
+                cache_dir='/tmp/models'
             )
             
             if torch.cuda.is_available():
                 self.inpaint_model.to("cuda")
-                self.inpaint_model.enable_attention_slicing()  # Reduce memory usage
+                self.inpaint_model.enable_attention_slicing()
             
             # 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
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    cache_dir='/tmp/models'
+                )
+                self.processor = SegformerImageProcessor.from_pretrained(
+                    "mattmdjaga/segformer_b2_clothes",
+                    cache_dir='/tmp/models'
                 )
-                self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
             except Exception as e:
                 print(f"Segmentation model loading error: {str(e)}")
                 raise
@@ -42,112 +54,23 @@ class RafayyVirtualTryOn:
             print(f"Initialization error: {str(e)}")
             raise
 
-    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 get_clothing_mask(self, image):
-        """Safely extract clothing mask"""
-        try:
-            # 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")
-            
-            # 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()
-            
-            # 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)
-            
-            # 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)}")
+    # ... (rest of the code remains the same)
 
-    def try_on(self, image, prompt, style_strength=0.7, progress=gr.Progress()):
-        """Main try-on function with comprehensive error handling"""
+# Initialize model with error handling and retry mechanism
+def initialize_model(max_retries=3):
+    for attempt in range(max_retries):
         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)
-            
-            progress(0.3, desc="Detecting clothing...")
-            mask = self.get_clothing_mask(original_image)
-            
-            # Clear GPU memory
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-                gc.collect()
-            
-            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(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
-            
+            return RafayyVirtualTryOn()
         except Exception as e:
-            error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
-            print(error_msg)  # For logging
-            raise gr.Error(str(e))
+            if attempt == max_retries - 1:
+                print(f"Failed to initialize model after {max_retries} attempts: {str(e)}")
+                raise
+            print(f"Attempt {attempt + 1} failed, retrying...")
+            torch.cuda.empty_cache()
+            gc.collect()
 
-# Initialize model with error handling
 try:
-    model = RafayyVirtualTryOn()
+    model = initialize_model()
 except Exception as e:
     print(f"Model initialization failed: {str(e)}")
     raise
@@ -186,7 +109,7 @@ demo = gr.Interface(
     cache_examples=True
 )
 
-# Launch with error handling
+# Launch with error handling and retry mechanism
 if __name__ == "__main__":
     try:
         demo.launch(
@@ -194,7 +117,9 @@ if __name__ == "__main__":
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            enable_queue=True
+            enable_queue=True,
+            cache_examples=True,
+            max_threads=4
         )
     except Exception as e:
         print(f"Launch failed: {str(e)}")