# 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()