|
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation |
|
from PIL import Image |
|
import requests |
|
import matplotlib.pyplot as plt |
|
import torch.nn as nn |
|
import torch |
|
from torchvision import transforms |
|
from transformers import SamModel, SamProcessor |
|
from diffusers import AutoPipelineForInpainting |
|
from diffusers.utils import load_image, make_image_grid |
|
|
|
|
|
|
|
def modify_image(image_url, prompt, mask_id=4): |
|
|
|
processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer_b3_clothes") |
|
model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer_b3_clothes") |
|
|
|
image = Image.open(image_url) |
|
inputs = processor(images=image, return_tensors="pt") |
|
|
|
outputs = model(**inputs) |
|
logits = outputs.logits.cpu() |
|
|
|
upsampled_logits = nn.functional.interpolate( |
|
logits, |
|
size=image.size[::-1], |
|
mode="bilinear", |
|
align_corners=False, |
|
) |
|
|
|
pred_seg = upsampled_logits.argmax(dim=1)[0] |
|
|
|
mask = (pred_seg == mask_id).numpy() |
|
mask_image = Image.fromarray((mask * 255).astype('uint8')) |
|
|
|
pipeline = AutoPipelineForInpainting.from_pretrained( |
|
"redstonehero/ReV_Animated_Inpainting", |
|
torch_dtype=torch.float16) |
|
|
|
|
|
|
|
image1 = pipeline(prompt=prompt, |
|
num_inference_steps=24, |
|
image=image, |
|
mask_image=mask_image, |
|
guidance_scale=3, |
|
strength=1.0).images[0] |
|
|
|
return make_image_grid([image1], rows = 1, cols = 1) |
|
|
|
|
|
|
|
import gradio as gr |
|
|
|
def gradio_wrapper(image, prompt, choice): |
|
return modify_image(image, prompt, int(choice)) |
|
|
|
demo = gr.Interface( |
|
fn=gradio_wrapper, |
|
inputs=[ |
|
gr.Image(type="filepath"), |
|
gr.Textbox(label="Prompt"), |
|
gr.Radio(["4", "6"], label="Mask ID") |
|
], |
|
outputs=gr.Image() |
|
) |
|
|
|
demo.launch(inline=False) |
