run model

app.py

@@ -1,7 +1,67 @@
 import gradio as gr
+from diffusers import DiffusionPipeline
+import torch
+from PIL import Image, ImageOps
+import requests
+from io import BytesIO
+from transparent_background import Remover
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Initialize the Diffusion Pipeline
+model_id = "yahoo-inc/photo-background-generation"
+pipeline = DiffusionPipeline.from_pretrained(model_id, custom_pipeline=model_id)
+pipeline = pipeline.to('cuda')
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def resize_with_padding(img, expected_size):
+    img.thumbnail((expected_size[0], expected_size[1]))
+    delta_width = expected_size[0] - img.size[0]
+    delta_height = expected_size[1] - img.size[1]
+    pad_width = delta_width // 2
+    pad_height = delta_height // 2
+    padding = (pad_width, pad_height, delta_width - pad_width, delta_height - pad_height)
+    return ImageOps.expand(img, padding)
+
+def process_image(input_image, prompt):
+    # Resize and process the input image
+    img = resize_with_padding(input_image, (512, 512))
+
+    # Load background detection model
+    remover = Remover(mode='base')
+
+    # Get foreground mask
+    fg_mask = remover.process(img, type='map')
+
+    seed = 13
+    mask = ImageOps.invert(fg_mask)
+    img = resize_with_padding(img, (512, 512))
+    generator = torch.Generator(device='cuda').manual_seed(seed)
+    cond_scale = 1.0
+
+    with torch.autocast("cuda"):
+        controlnet_image = pipeline(
+            prompt=prompt, 
+            image=img, 
+            mask_image=mask, 
+            control_image=mask, 
+            num_images_per_prompt=1, 
+            generator=generator, 
+            num_inference_steps=20, 
+            guess_mode=False, 
+            controlnet_conditioning_scale=cond_scale
+        ).images[0]
+
+    return controlnet_image
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=process_image, 
+    inputs=[
+        gr.inputs.Image(type="pil", label="Upload Image"), 
+        gr.inputs.Textbox(label="Enter Prompt")
+    ],
+    outputs=gr.outputs.Image(label="Generated Image"),
+    title="Image Processing with Diffusion Pipeline",
+    description="Upload an image and enter a prompt to generate a new image using the diffusion model."
+)
+
+# Launch the interface
+iface.launch()