Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from diffusers import StableDiffusionInpaintPipeline | |
| import torch | |
| from PIL import Image | |
| import numpy as np | |
| # Load the Stable Diffusion inpainting model (ensure GPU usage if available) | |
| pipe = StableDiffusionInpaintPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-2-1", # You can use another inpainting model here | |
| torch_dtype=torch.float16 | |
| ).to("cuda") | |
| # Function to remove makeup based on prompt and image | |
| def remove_makeup(image, prompt): | |
| # Debugging: Check the type of the input image | |
| print(f"Input image type: {type(image)}") | |
| # Ensure the input image is a PIL Image | |
| if not isinstance(image, Image.Image): | |
| image = Image.fromarray(image) | |
| # Debugging: Print the size before resizing | |
| print(f"Original image size: {image.size}") | |
| # Resize the image to 512x512 (standard input size for Stable Diffusion) | |
| image = image.resize((512, 512)) | |
| # Debugging: Print the size after resizing | |
| print(f"Resized image size: {image.size}") | |
| # Convert the image to a numpy array and normalize to range [-1, 1] | |
| image_array = np.array(image) / 255.0 | |
| image_array = (image_array - 0.5) / 0.5 # Normalize to [-1, 1] | |
| # Debugging: Check the shape and type of the image array | |
| print(f"Image array shape: {image_array.shape}, type: {type(image_array)}") | |
| # Convert to tensor and move to GPU | |
| image_tensor = torch.tensor(image_array).permute(2, 0, 1).unsqueeze(0).float().to("cuda") # Convert to tensor | |
| # Debugging: Check the shape and type of the tensor before sending it to the model | |
| print(f"Image tensor shape: {image_tensor.shape}, type: {type(image_tensor)}") | |
| # Use the model to process the image with the given prompt | |
| result = pipe(prompt=prompt, image=image_tensor).images[0] | |
| # Return the result | |
| return result | |
| # Define the Gradio interface | |
| iface = gr.Interface( | |
| fn=remove_makeup, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image with Makeup"), # Input image (PIL) | |
| gr.Textbox(label="Prompt", placeholder="Describe what you want (e.g., 'a natural face with no makeup and no filter')") # Input prompt | |
| ], | |
| outputs=gr.Image(type="pil", label="Makeup Removed Image"), | |
| title="Makeup Removal Using Stable Diffusion", | |
| description="Upload a face image with makeup and provide a prompt (e.g., 'a natural face with no makeup and no filter'). The model will remove makeup and filters, returning a natural-looking face.", | |
| ) | |
| # Launch the Gradio interface | |
| iface.launch() | |