Implement outpainting pipeline and update related files

runner/Dockerfile

@@ -54,4 +54,6 @@ COPY images/ /app/images
 COPY bench.py /app/bench.py
 COPY example_data/ /app/example_data
 
-CMD ["uvicorn", "app.main:app", "--log-config", "app/cfg/uvicorn_logging_config.json", "--host", "0.0.0.0", "--port", "8000"]
+EXPOSE 8700
+
+CMD ["uvicorn", "app.main:app", "--log-config", "app/cfg/uvicorn_logging_config.json", "--host", "0.0.0.0", "--port", "8700"]

runner/app/main.py

@@ -5,6 +5,9 @@
 from app.routes import health
 from fastapi import FastAPI
 from fastapi.routing import APIRoute
+from .routes import image_outpainting  
+from .pipelines.image_outpainting import ImageOutpaintingPipeline
+from app.routes.image_outpainting import router as image_router
 
 logger = logging.getLogger(__name__)
 
@@ -57,6 +60,8 @@ def load_pipeline(pipeline: str, model_id: str) -> any:
         case "llm":
             from app.pipelines.llm import LLMPipeline
             return LLMPipeline(model_id)
+        case "image-outpainting":
+            return ImageOutpaintingPipeline(model_id)
         case _:
             raise EnvironmentError(
                 f"{pipeline} is not a valid pipeline for model {model_id}"
@@ -94,6 +99,8 @@ def load_route(pipeline: str) -> any:
         case "llm":
             from app.routes import llm
             return llm.router
+        case "image-outpainting":
+            return image_router
         case _:
             raise EnvironmentError(f"{pipeline} is not a valid pipeline")
 
@@ -113,3 +120,4 @@ def use_route_names_as_operation_ids(app: FastAPI) -> None:
 
 
 app = FastAPI(lifespan=lifespan)
+app.include_router(image_router)

runner/app/pipelines/image_outpainting.py

@@ -0,0 +1,54 @@
+from diffusers import AutoPipelineForInpainting
+import torch
+from PIL import Image
+import numpy as np
+
+class ImageOutpaintingPipeline:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        try:
+            # Use AutoPipelineForInpainting to load ProPainter
+            self.pipe = AutoPipelineForInpainting.from_pretrained("ruffy369/propainter", torch_dtype=torch.float16).to(self.device)
+            print("ProPainter model loaded successfully.")
+        except Exception as e:
+            print(f"Error loading ProPainter model: {e}")
+            self.pipe = None
+
+    def __call__(
+        self,
+        image: Image.Image,
+        prompt: str,
+        negative_prompt: str = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+    ):
+        if self.pipe is None:
+            print("ProPainter model is not loaded. Cannot perform outpainting.")
+            return None
+
+        # Prepare the image for outpainting
+        width, height = image.size
+        target_size = min(max(width, height) * 2, 1024)  # Double the size, but cap at 1024
+        new_image = Image.new('RGB', (target_size, target_size), (255, 255, 255))
+        new_image.paste(image, ((target_size - width) // 2, (target_size - height) // 2))
+
+        # Create a mask for outpainting
+        mask = Image.new('L', (target_size, target_size), 255)
+        mask.paste(0, ((target_size - width) // 2, (target_size - height) // 2, 
+                       (target_size + width) // 2, (target_size + height) // 2))
+
+        try:
+            # Generate the outpainted image
+            output = self.pipe(
+                prompt=prompt,
+                image=new_image,
+                mask_image=mask,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+            ).images[0]
+
+            return output
+        except Exception as e:
+            print(f"Error during outpainting: {e}")
+            return None

runner/app/routes/image_outpainting.py

@@ -0,0 +1,57 @@
+from fastapi import APIRouter, File, UploadFile, Form, HTTPException, Query
+from PIL import Image
+import io
+from ..pipelines.image_outpainting import ImageOutpaintingPipeline
+from .util import ImageOutpaintingResponse
+
+
+router = APIRouter()
+pipeline = ImageOutpaintingPipeline()
+
+
+
+def resize_image(image: Image.Image, max_size: int = 1024) -> Image.Image:
+    """Resize image while maintaining aspect ratio if it exceeds max_size"""
+    if max(image.size) > max_size:
+        image.thumbnail((max_size, max_size))
+    return image
+
+@router.post("/out-paint", response_model=ImageOutpaintingResponse)
+async def out_paint(
+    image: UploadFile = File(...),
+    prompt: str = Form(...),
+    negative_prompt: str = Form(None),
+    num_inference_steps: int = Form(50, ge=1, le=1000),
+    guidance_scale: float = Form(7.5, ge=0, le=20),
+):
+    if len(prompt) > 1000:
+        raise HTTPException(status_code=400, detail="Prompt is too long")
+    if negative_prompt and len(negative_prompt) > 1000:
+        raise HTTPException(status_code=400, detail="Negative prompt is too long")
+
+    try:
+        image_content = await image.read()
+        input_image = resize_image(Image.open(io.BytesIO(image_content)).convert("RGB"))
+
+        output_image = pipeline(
+            image=input_image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        )
+
+        # Convert the output image to bytes for response
+        img_byte_arr = io.BytesIO()
+        output_image.save(img_byte_arr, format='PNG')
+        img_byte_arr = img_byte_arr.getvalue()
+
+        return ImageOutpaintingResponse(
+            image=img_byte_arr,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred during outpainting: {str(e)}")

runner/app/routes/util.py

@@ -181,3 +181,15 @@ def json_str_to_np_array(
             error_message += f": {e}"
             raise ValueError(error_message)
     return None
+
+
+class ImageOutpaintingResponse(BaseModel):
+    """
+    Response model for the image outpainting operation.
+    Contains the resulting outpainted image and the parameters used in the process."""
+
+    image: bytes = Field(..., description="The outpainted image in bytes format.")
+    prompt: str = Field(..., description="The prompt used for outpainting.")
+    negative_prompt: str = Field(None, description="The negative prompt used for outpainting, if any.")
+    num_inference_steps: int = Field(..., description="The number of inference steps used.")
+    guidance_scale: float = Field(..., description="The guidance scale used for outpainting.")

runner/dl_checkpoints.sh

@@ -43,6 +43,16 @@ function download_beta_models() {
     # Download image-to-video models (token-gated).
     check_hf_auth
     huggingface-cli download stabilityai/stable-video-diffusion-img2vid-xt-1-1 --include "*.fp16.safetensors" "*.json" --cache-dir models ${TOKEN_FLAG:+"$TOKEN_FLAG"}
+
+    # Download ProPainter and stable-outpainting models
+    printf "\nDownloading outpainting models...\n"
+    huggingface-cli download ruffy369/propainter --cache-dir models
+    # Only download stable-outpainting if ProPainter fails
+    if [ $? -ne 0 ]; then
+        printf "Failed to download ProPainter model. Downloading stable-outpainting as backup...\n"
+        huggingface-cli download Brvcket/stable-outpainting-xl-0.1 --cache-dir models
+    fi
+
 }
 
 # Download all models.