Update app.py

app.py

@@ -1,60 +1,43 @@
 import gradio as gr
-from gradio_imageslider import ImageSlider
-from loadimg import load_img
-import spaces
-from transformers import AutoModelForImageSegmentation
 import torch
-from torchvision import transforms
-
-torch.set_float32_matmul_precision(["high", "highest"][0])
-
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
-birefnet.to("cuda")
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
-
-
-@spaces.GPU
-def fn(image):
-    im = load_img(image, output_type="pil")
-    im = im.convert("RGB")
-    image_size = im.size
-    origin = im.copy()
-    image = load_img(im)
-    input_images = transform_image(image).unsqueeze(0).to("cuda")
-    # Prediction
-    with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-    image.putalpha(mask)
-    return (image, origin)
-
-
-slider1 = ImageSlider(label="birefnet", type="pil")
-slider2 = ImageSlider(label="birefnet", type="pil")
-image = gr.Image(label="Upload an image")
-text = gr.Textbox(label="Paste an image URL")
-
-
-chameleon = load_img("chameleon.jpg", output_type="pil")
-
-url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
-tab1 = gr.Interface(
-    fn, inputs=image, outputs=slider1, examples=[chameleon], api_name="image"
-)
-
-tab2 = gr.Interface(fn, inputs=text, outputs=slider2, examples=[url], api_name="text")
-
-
-demo = gr.TabbedInterface(
-    [tab1, tab2], ["image", "text"], title="birefnet for background removal"
-)
+from diffusers import StableDiffusionInpaintPipeline
+
+# Load the Stable Diffusion inpainting model
+def load_inpainting_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
+    model.to(device)
+    return model
+
+# Function to edit the clothing image based on prompt
+def edit_clothing_image(prompt, image):
+    # Load the model
+    model = load_inpainting_model()
+    
+    # Inpainting to edit the image based on the text prompt
+    edited_image = model(prompt=prompt, image=image).images[0]
+    return edited_image
+
+# Gradio interface
+def interface():
+    with gr.Blocks() as ui:
+        gr.Markdown("# Clothing Image Editing with AI")
+
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(label="Editing Prompt", placeholder="e.g., change the color of the t-shirt to red")
+                input_image = gr.Image(label="Upload Image", tool="editor", type="pil")
+                output_image = gr.Image(label="Edited Image")
+
+                edit_button = gr.Button("Apply Edit")
+
+        edit_button.click(edit_clothing_image, inputs=[prompt, input_image], outputs=output_image)
+
+    return ui
+
+# Launch the interface
+ui = interface()
+ui.launch()