app.py

# app.py - Main application file for Hugging Face Spaces
import gradio as gr
import torch
import torch.nn.functional as F
from diffusers import FluxPipeline
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
import numpy as np
from PIL import Image
import cv2
import os
from typing import Optional, Tuple
import warnings
warnings.filterwarnings("ignore")

class ClothingSegmentationModel:
    def __init__(self, model_name="mattmdjaga/segformer_b2_clothes"):
        self.feature_extractor = SegformerFeatureExtractor.from_pretrained(model_name)
        self.model = SegformerForSemanticSegmentation.from_pretrained(model_name)
        self.model.eval()
        
    def segment_clothing(self, image: Image.Image, target_classes: list = None) -> np.ndarray:
        if target_classes is None:
            target_classes = [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
        
        inputs = self.feature_extractor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            
        logits = F.interpolate(logits, size=image.size[::-1], mode='bilinear', align_corners=False)
        predicted_mask = torch.argmax(logits, dim=1)[0].cpu().numpy()
        
        clothing_mask = np.zeros_like(predicted_mask)
        for class_id in target_classes:
            clothing_mask[predicted_mask == class_id] = 1
            
        return clothing_mask

class FluxTryOffModel:
    def __init__(self, model_path: str = "black-forest-labs/FLUX.1-dev"):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Initialize the pipeline
        self.pipe = FluxPipeline.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None
        )
        
        if torch.cuda.is_available():
            self.pipe.to(self.device)
        
        self.segmentation_model = ClothingSegmentationModel()
        
    def preprocess_image(self, image: Image.Image, target_size: Tuple[int, int] = (512, 512)) -> Image.Image:
        if image.mode != 'RGB':
            image = image.convert('RGB')
            
        image.thumbnail(target_size, Image.Resampling.LANCZOS)
        
        width, height = image.size
        target_width, target_height = target_size
        
        if width != target_width or height != target_height:
            new_image = Image.new('RGB', target_size, (255, 255, 255))
            left = (target_width - width) // 2
            top = (target_height - height) // 2
            new_image.paste(image, (left, top))
            image = new_image
            
        return image
    
    def create_inpainting_mask(self, image: Image.Image, user_mask: Optional[np.ndarray] = None) -> np.ndarray:
        if user_mask is not None:
            return user_mask
        else:
            return self.segmentation_model.segment_clothing(image)
    
    def generate_tryoff_result(
        self, 
        image: Image.Image, 
        mask: np.ndarray,
        prompt: str = "natural skin texture, realistic human body, high quality",
        negative_prompt: str = "clothing, fabric, shirt, dress, blurry, low quality",
        num_inference_steps: int = 20,
        guidance_scale: float = 7.5,
        strength: float = 0.85
    ) -> Image.Image:
        processed_image = self.preprocess_image(image)
        
        mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
        mask_pil = mask_pil.resize(processed_image.size, Image.Resampling.NEAREST)
        
        result = self.pipe(
            prompt=prompt,
            image=processed_image,
            mask_image=mask_pil,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            strength=strength
        ).images[0]
        
        return result

# Global model instance
model = None

def initialize_model():
    global model
    if model is None:
        model = FluxTryOffModel()
    return model

def process_image(
    input_image: Image.Image,
    mask_image: Optional[Image.Image] = None,
    prompt: str = "natural skin texture, realistic human body, high quality",
    negative_prompt: str = "clothing, fabric, shirt, dress, blurry, low quality",
    num_steps: int = 20,
    guidance_scale: float = 7.5,
    strength: float = 0.85
):
    if input_image is None:
        return None, None, "Please upload an image"
    
    try:
        # Initialize model if not already done
        model = initialize_model()
        
        # Create mask
        if mask_image is not None:
            mask_array = np.array(mask_image.convert('L'))
            mask_array = (mask_array > 128).astype(np.uint8)
        else:
            mask_array = None
        
        # Create inpainting mask
        final_mask = model.create_inpainting_mask(input_image, mask_array)
        
        # Generate result
        result = model.generate_tryoff_result(
            image=input_image,
            mask=final_mask,
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            strength=strength
        )
        
        # Create visualization of the mask
        mask_vis = Image.fromarray((final_mask * 255).astype(np.uint8))
        
        return result, mask_vis, "Success!"
        
    except Exception as e:
        return None, None, f"Error: {str(e)}"

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Virtual Try-Off Model", theme=gr.themes.Soft()) as interface:
        gr.Markdown("# 🎭 Virtual Try-Off Model")
        gr.Markdown("Upload an image and optionally draw a mask to remove clothing. The model will generate a realistic result with the clothing removed.")
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="📸 Input Image",
                    type="pil",
                    height=400
                )
                
                mask_image = gr.Image(
                    label="🎨 Mask (Optional - draw on clothing to remove)",
                    type="pil",
                    height=300,
                    tool="sketch"
                )
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    prompt = gr.Textbox(
                        label="Prompt",
                        value="natural skin texture, realistic human body, high quality",
                        lines=2
                    )
                    
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        value="clothing, fabric, shirt, dress, blurry, low quality",
                        lines=2
                    )
                    
                    num_steps = gr.Slider(
                        label="Number of Steps",
                        minimum=10,
                        maximum=50,
                        value=20,
                        step=1
                    )
                    
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=1.0,
                        maximum=20.0,
                        value=7.5,
                        step=0.1
                    )
                    
                    strength = gr.Slider(
                        label="Strength",
                        minimum=0.1,
                        maximum=1.0,
                        value=0.85,
                        step=0.05
                    )
            
            with gr.Column():
                output_image = gr.Image(
                    label="✨ Try-Off Result",
                    type="pil",
                    height=400
                )
                
                mask_visualization = gr.Image(
                    label="🎯 Generated Mask",
                    type="pil",
                    height=300
                )
                
                status = gr.Textbox(
                    label="Status",
                    value="Ready",
                    interactive=False
                )
        
        process_btn = gr.Button("🚀 Generate Try-Off Result", variant="primary", size="lg")
        
        process_btn.click(
            fn=process_image,
            inputs=[
                input_image,
                mask_image,
                prompt,
                negative_prompt,
                num_steps,
                guidance_scale,
                strength
            ],
            outputs=[output_image, mask_visualization, status]
        )
        
        # Add usage instructions
        gr.Markdown("""
        ## 📋 Instructions:
        1. Upload an image of a person wearing clothing
        2. (Optional) Draw a mask over the clothing area you want to remove
        3. Adjust the advanced settings if needed
        4. Click "Generate Try-Off Result" to process
        
        ## ⚠️ Notes:
        - The model works best with clear, front-facing images
        - Processing may take 30-60 seconds depending on settings
        - Higher steps and guidance scale = better quality but slower processing
        """)
    
    return interface

# Launch the app
if __name__ == "__main__":
    app = create_interface()
    app.launch()