Initial commit

.gitignore

@@ -0,0 +1,178 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+.env
+.venv
+venv/
+ENV/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+.DS_Store
+
+# Binary files and assets
+parser/u2net_cloth_seg/assets/*.png
+upscaler/real_esrgan/assets/*.png
+upscaler/real_esrgan/assets/*.jpg
+upscaler/real_esrgan/inputs/*.png
+upscaler/real_esrgan/inputs/video/*.mp4
+upscaler/real_esrgan/tests/data/gt.lmdb/
+upscaler/real_esrgan/tests/data/gt/*.png
+
+# Models
+models/
+*.pth
+*.ckpt
+*.safetensors
+
+# Logs
+*.log
+logs/
+
+# Temporary files
+*.tmp
+*.temp
+*.bak
+*.backup
+
+# System files
+.DS_Store
+Thumbs.db
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Additional exclusions
+*.swp
+*.swo

Dockerfile

@@ -0,0 +1,41 @@
+# Must use a Cuda version 11+
+# FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
+FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
+
+WORKDIR /
+COPY ./parser /parser
+COPY ./configs /configs
+RUN mkdir /checkpoints
+# Install git
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    build-essential
+
+# Install python packages
+RUN pip3 install --upgrade pip
+ADD requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+# Install cv2 dependencies
+RUN apt-get update
+RUN apt-get install ffmpeg libsm6 libxext6  -y
+
+# We add the banana boilerplate here
+ADD server.py .
+
+# Add your model weight files 
+# (in this case we have a python script)
+ADD download.py .
+
+# Add your custom app code, init() and inference()
+ADD app.py .
+
+ENV PYTHONPATH "${PYTHONPATH}:/parser:/upscaler"
+
+#Alternative to using build args, you can put your token in the next line
+#ENV HF_AUTH_TOKEN={token}
+RUN python3 download.py
+
+EXPOSE 8000
+
+CMD python3 -u server.py

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Banana
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,38 @@
+---
+title: ClothQuill - AI Clothing Inpainting
+emoji: 👕
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 5.25.1
+app_file: app.py
+pinned: false
+---
+
+# ClothQuill - AI Clothing Inpainting
+
+This Space allows you to inpaint clothing on images using Stable Diffusion. Upload an image, provide a prompt describing the clothing you want to generate, and get multiple inpainted results.
+
+## How to Use
+
+1. Upload an image containing a person
+2. Enter a prompt describing the clothing you want to generate
+3. Click "Generate" to get multiple inpainted results
+4. Download your favorite result
+
+## Examples
+
+- Prompt: "A stylish black leather jacket"
+- Prompt: "A formal blue suit with white shirt"
+- Prompt: "A casual red hoodie"
+
+## Technical Details
+
+This Space uses:
+- Stable Diffusion for inpainting
+- U2NET for human parsing
+- RealESRGAN for upscaling
+
+## License
+
+This project is licensed under the MIT License.

app.py

@@ -0,0 +1,429 @@
+import torch
+from torch import autocast
+from diffusers import StableDiffusionInpaintPipeline
+import gradio as gr
+import traceback
+import base64
+from io import BytesIO
+import os
+# import sys
+import PIL
+import json
+import requests
+import logging
+import time
+import warnings
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+warnings.filterwarnings("ignore")
+
+# sys.path.insert(1, './parser')
+
+# from parser.schp_masker import *
+from parser.segformer_parser import SegformerParser
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger('clothquill')
+
+# Model paths
+SEGFORMER_MODEL = "mattmdjaga/segformer_b2_clothes"
+STABLE_DIFFUSION_MODEL = "stabilityai/stable-diffusion-2-inpainting"
+
+# Global variables for models
+parser = None
+model = None
+inpainter = None
+original_image = None  # Store the original uploaded image
+
+# Color mapping for different clothing parts
+CLOTHING_COLORS = {
+    'Background': (0, 0, 0, 0),      # Transparent
+    'Hat': (255, 0, 0, 128),         # Red
+    'Hair': (0, 255, 0, 128),        # Green
+    'Glove': (0, 0, 255, 128),       # Blue
+    'Sunglasses': (255, 255, 0, 128), # Yellow
+    'Upper-clothes': (255, 0, 255, 128), # Magenta
+    'Dress': (0, 255, 255, 128),     # Cyan
+    'Coat': (128, 0, 0, 128),        # Dark Red
+    'Socks': (0, 128, 0, 128),       # Dark Green
+    'Pants': (0, 0, 128, 128),       # Dark Blue
+    'Jumpsuits': (128, 128, 0, 128), # Dark Yellow
+    'Scarf': (128, 0, 128, 128),     # Dark Magenta
+    'Skirt': (0, 128, 128, 128),     # Dark Cyan
+    'Face': (192, 192, 192, 128),    # Light Gray
+    'Left-arm': (64, 64, 64, 128),   # Dark Gray
+    'Right-arm': (64, 64, 64, 128),  # Dark Gray
+    'Left-leg': (32, 32, 32, 128),   # Very Dark Gray
+    'Right-leg': (32, 32, 32, 128),  # Very Dark Gray
+    'Left-shoe': (16, 16, 16, 128),  # Almost Black
+    'Right-shoe': (16, 16, 16, 128), # Almost Black
+}
+
+def get_device():
+    if torch.cuda.is_available():
+        device = "cuda"
+        logger.info("Using GPU")
+    else:
+        device = "cpu"
+        logger.info("Using CPU")
+    return device
+
+def init():
+    global parser
+    global model
+    global inpainter
+
+    start_time = time.time()
+    logger.info("Starting application initialization")
+
+    try:
+        device = get_device()
+        
+        # Check if models directory exists
+        if not os.path.exists("models"):
+            logger.info("Creating models directory...")
+            from download_models import download_models
+            download_models()
+
+        # Initialize Segformer parser
+        logger.info("Initializing Segformer parser...")
+        parser = SegformerParser(SEGFORMER_MODEL)
+        
+        # Initialize Stable Diffusion model
+        logger.info("Initializing Stable Diffusion model...")
+        model = StableDiffusionInpaintPipeline.from_pretrained(
+            STABLE_DIFFUSION_MODEL,
+            safety_checker=None,
+            revision="fp16" if device == "cuda" else None,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        ).to(device)
+        
+        # Initialize inpainter
+        logger.info("Initializing inpainter...")
+        inpainter = ClothingInpainter(model=model, parser=parser)
+        
+        logger.info(f"Application initialized in {time.time() - start_time:.2f} seconds")
+    except Exception as e:
+        logger.error(f"Error initializing application: {str(e)}")
+        raise e
+
+class ClothingInpainter:
+    def __init__(self, model_path=None, model=None, parser=None):
+        self.device = get_device()
+        self.last_mask = None  # Store the last generated mask
+        self.original_image = None  # Store the original image
+        
+        if model_path is None and model is None:
+            raise ValueError('No model provided!')
+        if model_path is not None:
+            self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                model_path,
+                safety_checker=None,
+                revision="fp16" if self.device == "cuda" else None,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+            ).to(self.device)
+        else:
+            self.pipe = model
+
+        self.parser = parser
+
+    def make_square(self, im, min_size=256, fill_color=(0, 0, 0, 0)):
+        x, y = im.size
+        size = max(min_size, x, y)
+        new_im = PIL.Image.new('RGBA', (size, size), fill_color)
+        new_im.paste(im, (int((size - x) / 2), int((size - y) / 2)))
+        return new_im.convert('RGB')
+
+    def unmake_square(self, init_im, op_im, min_size=256, rs_size=512):
+        x, y = init_im.size
+        size = max(min_size, x, y)
+        factor = rs_size/size
+        return op_im.crop((int((size-x) * factor / 2), int((size-y) * factor / 2),\
+                                int((size+x) * factor / 2), int((size+y) * factor / 2)))
+
+    def visualize_segmentation(self, image, masks, selected_parts=None):
+        """Visualize segmentation with colored overlays for selected parts and gray for unselected."""
+        # Always use original image if available
+        image_to_use = self.original_image if self.original_image is not None else image
+        
+        # Create a copy of the original image
+        original_size = image_to_use.size
+        vis_image = image_to_use.copy().convert('RGBA')
+        
+        # Create overlay at 512x512
+        overlay = Image.new('RGBA', (512, 512), (0, 0, 0, 0))
+        draw = ImageDraw.Draw(overlay)
+        
+        # Draw each mask with its corresponding color
+        for part_name, mask in masks.items():
+            # Convert part name for color lookup
+            color_key = part_name.replace('-', ' ').title().replace(' ', '-')
+            is_selected = selected_parts and part_name in selected_parts
+            
+            # If selected, use color (with fallback). If unselected, use faint gray
+            if is_selected:
+                color = CLOTHING_COLORS.get(color_key, (255, 0, 255, 128))  # Default to magenta if no color found
+            else:
+                color = (180, 180, 180, 80)  # Faint gray for unselected
+            
+            mask_array = np.array(mask)
+            coords = np.where(mask_array > 0)
+            for y, x in zip(coords[0], coords[1]):
+                draw.point((x, y), fill=color)
+        
+        # Resize overlay to match original image size
+        overlay = overlay.resize(original_size, Image.Resampling.LANCZOS)
+        
+        # Composite the overlay onto the original image
+        vis_image = Image.alpha_composite(vis_image, overlay)
+        return vis_image
+
+    def inpaint(self, prompt, init_image, selected_parts=None, dilation_iterations=2) -> dict:
+        image = self.make_square(init_image).resize((512,512))
+
+        if self.parser is not None:
+            masks = self.parser.get_all_masks(image)
+            masks = {k: v.resize((512,512)) for k, v in masks.items()}
+        else:
+            raise ValueError('Image Parser is Missing')
+        
+        logger.info(f'[generated required mask(s) at {time.time()}]')
+
+        # Create combined mask for selected parts
+        if selected_parts:
+            combined_mask = Image.new('L', (512, 512), 0)
+            for part in selected_parts:
+                if part in masks:
+                    mask_array = np.array(masks[part])
+                    kernel = np.ones((5,5), np.uint8)
+                    dilated_mask = cv2.dilate(mask_array, kernel, iterations=dilation_iterations)
+                    dilated_mask = Image.fromarray(dilated_mask)
+                    combined_mask = Image.composite(
+                        Image.new('L', (512, 512), 255),
+                        combined_mask,
+                        dilated_mask
+                    )
+        else:
+            # If no parts selected, use all clothing parts
+            combined_mask = Image.new('L', (512, 512), 0)
+            for part, mask in masks.items():
+                if part in ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']:
+                    mask_array = np.array(mask)
+                    kernel = np.ones((5,5), np.uint8)
+                    dilated_mask = cv2.dilate(mask_array, kernel, iterations=dilation_iterations)
+                    dilated_mask = Image.fromarray(dilated_mask)
+                    combined_mask = Image.composite(
+                        Image.new('L', (512, 512), 255),
+                        combined_mask,
+                        dilated_mask
+                    )
+
+        # Run the model
+        guidance_scale=7.5
+        num_samples = 3
+        with autocast("cuda"), torch.inference_mode():
+            images = self.pipe(
+                num_inference_steps = 50,
+                prompt=prompt['pos'],
+                image=image,
+                mask_image=combined_mask,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=num_samples,
+            ).images
+        
+        images_output = []
+        for img in images:
+            ch = PIL.Image.composite(img, image, combined_mask)
+            fin_img = self.unmake_square(init_image, ch)
+            images_output.append(fin_img)
+
+        return images_output
+
+def process_segmentation(image, dilation_iterations=2):
+    try:
+        if image is None:
+            raise gr.Error("Please upload an image")
+            
+        # Store original image
+        inpainter.original_image = image.copy()
+            
+        # Create a processing copy at 512x512
+        proc_image = image.resize((512, 512), Image.Resampling.LANCZOS)
+            
+        # Get the main mask
+        all_masks = inpainter.parser.get_all_masks(proc_image)
+        if not all_masks:
+            logger.error("No clothing detected in the image")
+            raise gr.Error("No clothing detected in the image. Please try a different image.")
+        inpainter.last_mask = all_masks
+        # Only show main clothing parts for selection
+        main_parts = ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']
+        masks = {k: v for k, v in all_masks.items() if k in main_parts}
+        vis_image = inpainter.visualize_segmentation(image, masks, selected_parts=None)
+        detected_parts = [k for k in masks.keys()]
+        return vis_image, gr.update(choices=detected_parts, value=[])
+    except gr.Error as e:
+        raise e
+    except Exception as e:
+        logger.error(f"Error processing segmentation: {str(e)}")
+        raise gr.Error("Error processing the image. Please try a different image.")
+
+def update_dilation(image, selected_parts, dilation_iterations):
+    try:
+        if image is None or inpainter.last_mask is None:
+            return image
+        # Redilate all stored masks
+        main_parts = ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']
+        masks = {}
+        for part in main_parts:
+            if part in inpainter.last_mask:
+                mask_array = np.array(inpainter.last_mask[part])
+                kernel = np.ones((5,5), np.uint8)
+                dilated_mask = cv2.dilate(mask_array, kernel, iterations=dilation_iterations)
+                masks[part] = Image.fromarray(dilated_mask)
+        # Use original image for visualization
+        vis_image = inpainter.visualize_segmentation(inpainter.original_image, masks, selected_parts=selected_parts)
+        return vis_image
+    except Exception as e:
+        logger.error(f"Error updating dilation: {str(e)}")
+        return image
+
+def process_image(prompt, image, selected_parts, dilation_iterations):
+    start_time = time.time()
+    logger.info(f"Processing new request - Prompt: {prompt}, Image size: {image.size if image else 'None'}")
+    
+    try:
+        if image is None:
+            logger.error("No image provided")
+            raise gr.Error("Please upload an image")
+        if not prompt:
+            logger.error("No prompt provided")
+            raise gr.Error("Please enter a prompt")
+        if not selected_parts:
+            logger.error("No parts selected")
+            raise gr.Error("Please select at least one clothing part to modify")
+            
+        prompt_dict = {'pos': prompt}
+        logger.info("Starting inpainting process")
+        
+        # Generate inpainted images
+        # Convert selected_parts to lowercase/dash format
+        selected_parts = [p.lower() for p in selected_parts]
+        images = inpainter.inpaint(prompt_dict, image, selected_parts, dilation_iterations)
+        
+        if not images:
+            logger.error("Inpainting failed to produce results")
+            raise gr.Error("Failed to generate images. Please try again.")
+            
+        logger.info(f"Request processed in {time.time() - start_time:.2f} seconds")
+        return images
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise gr.Error(f"Error processing image: {str(e)}")
+
+def update_selected_parts(image, selected_parts, dilation_iterations):
+    try:
+        if image is None or inpainter.last_mask is None:
+            return image
+        main_parts = ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']
+        masks = {}
+        for part in main_parts:
+            if part in inpainter.last_mask:
+                mask_array = np.array(inpainter.last_mask[part])
+                kernel = np.ones((5,5), np.uint8)
+                dilated_mask = cv2.dilate(mask_array, kernel, iterations=dilation_iterations)
+                masks[part] = Image.fromarray(dilated_mask)
+        # Lowercase the selected_parts for comparison
+        selected_parts = [p.lower() for p in selected_parts] if selected_parts else []
+        # Use original image for visualization
+        vis_image = inpainter.visualize_segmentation(inpainter.original_image, masks, selected_parts=selected_parts)
+        return vis_image
+    except Exception as e:
+        logger.error(f"Error updating selected parts: {str(e)}")
+        return image
+
+# Initialize the model
+init()
+
+# Create Gradio interface
+with gr.Blocks(title="ClothQuill - AI Clothing Inpainting") as demo:
+    gr.Markdown("# ClothQuill - AI Clothing Inpainting")
+    gr.Markdown("Upload an image to see segmented clothing parts, then select parts to modify and describe your changes")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                type="pil", 
+                label="Upload Image",
+                scale=1,  # This ensures the image maintains its aspect ratio
+                height=None  # Allow dynamic height based on content
+            )
+            dilation_slider = gr.Slider(
+                minimum=0,
+                maximum=5,
+                value=2,
+                step=1,
+                label="Mask Dilation",
+                info="Adjust the mask dilation to control the area of modification"
+            )
+            selected_parts = gr.CheckboxGroup(
+                choices=[],
+                label="Select parts to modify",
+                value=[]
+            )
+            prompt = gr.Textbox(
+                label="Describe the clothing you want to generate",
+                placeholder="e.g., A stylish black leather jacket"
+            )
+            generate_btn = gr.Button("Generate")
+        
+        with gr.Column():
+            gallery = gr.Gallery(
+                label="Generated Results", 
+                show_label=False,
+                columns=2,
+                height=None,  # Allow dynamic height
+                object_fit="contain"  # Maintain aspect ratio
+            )
+    
+    # Add event handler for image upload
+    input_image.upload(
+        fn=process_segmentation,
+        inputs=[input_image, dilation_slider],
+        outputs=[input_image, selected_parts]
+    )
+    
+    # Add event handler for dilation changes
+    dilation_slider.change(
+        fn=update_dilation,
+        inputs=[input_image, selected_parts,dilation_slider],
+        outputs=input_image
+    )
+    
+    # Add event handler for generation
+    generate_btn.click(
+        fn=process_image,
+        inputs=[prompt, input_image, selected_parts, dilation_slider],
+        outputs=gallery
+    )
+
+    # Add event handler for part selection changes
+    selected_parts.change(
+        fn=update_selected_parts,
+        inputs=[input_image, selected_parts, dilation_slider],
+        outputs=input_image
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)
+
+    
+
+    
+
+

colab_demo.py

@@ -0,0 +1,201 @@
+import torch
+from torch import autocast
+from diffusers import StableDiffusionInpaintPipeline
+import gradio as gr
+import traceback
+import base64
+from io import BytesIO
+import os
+import PIL
+import json
+import requests
+import logging
+import time
+import warnings
+warnings.filterwarnings("ignore")
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger('looks.studio')
+
+# Model paths
+SEGFORMER_MODEL = "mattmdjaga/segformer_b2_clothes"
+STABLE_DIFFUSION_MODEL = "stabilityai/stable-diffusion-2-inpainting"
+
+# Global variables for models
+parser = None
+model = None
+inpainter = None
+
+def get_device():
+    if torch.cuda.is_available():
+        device = "cuda"
+        logger.info("Using GPU")
+    else:
+        device = "cpu"
+        logger.info("Using CPU")
+    return device
+
+def init():
+    global parser
+    global model
+    global inpainter
+
+    start_time = time.time()
+    logger.info("Starting application initialization")
+
+    try:
+        device = get_device()
+        
+        # Initialize Segformer parser
+        logger.info("Initializing Segformer parser...")
+        from parser.segformer_parser import SegformerParser
+        parser = SegformerParser(SEGFORMER_MODEL)
+        
+        # Initialize Stable Diffusion model
+        logger.info("Initializing Stable Diffusion model...")
+        model = StableDiffusionInpaintPipeline.from_pretrained(
+            STABLE_DIFFUSION_MODEL,
+            safety_checker=None,
+            revision="fp16" if device == "cuda" else None,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        ).to(device)
+        
+        # Initialize inpainter
+        logger.info("Initializing inpainter...")
+        inpainter = ClothingInpainter(model=model, parser=parser)
+        
+        logger.info(f"Application initialized in {time.time() - start_time:.2f} seconds")
+    except Exception as e:
+        logger.error(f"Error initializing application: {str(e)}")
+        raise e
+
+class ClothingInpainter:
+    def __init__(self, model_path=None, model=None, parser=None):
+        self.device = get_device()
+        
+        if model_path is None and model is None:
+            raise ValueError('No model provided!')
+        if model_path is not None:
+            self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                model_path,
+                safety_checker=None,
+                revision="fp16" if self.device == "cuda" else None,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+            ).to(self.device)
+        else:
+            self.pipe = model
+
+        self.parser = parser
+
+    def make_square(self, im, min_size=256, fill_color=(0, 0, 0, 0)):
+        x, y = im.size
+        size = max(min_size, x, y)
+        new_im = PIL.Image.new('RGBA', (size, size), fill_color)
+        new_im.paste(im, (int((size - x) / 2), int((size - y) / 2)))
+        return new_im.convert('RGB')
+
+    def unmake_square(self, init_im, op_im, min_size=256, rs_size=512):
+        x, y = init_im.size
+        size = max(min_size, x, y)
+        factor = rs_size/size
+        return op_im.crop((int((size-x) * factor / 2), int((size-y) * factor / 2),\
+                                int((size+x) * factor / 2), int((size+y) * factor / 2)))
+
+    def inpaint(self, prompt, init_image, parser=None) -> dict:
+        image = self.make_square(init_image).resize((512,512))
+
+        if self.parser is not None:
+            mask = self.parser.get_image_mask(image)
+            mask = mask.resize((512,512))
+        elif parser is not None:
+            mask = parser.get_image_mask(image)
+            mask = mask.resize((512,512))
+        else:
+            raise ValueError('Image Parser is Missing')
+        logger.info(f'[generated required mask(s) at {time.time()}]')
+
+        # Run the model
+        guidance_scale=7.5
+        num_samples = 3
+        with autocast("cuda"), torch.inference_mode():
+            images = self.pipe(
+                num_inference_steps = 50,
+                prompt=prompt['pos'],
+                image=image,
+                mask_image=mask,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=num_samples,
+            ).images
+        
+        images_output = []
+        for img in images:
+            ch = PIL.Image.composite(img,image, mask.convert('L'))
+            fin_img = self.unmake_square(init_image, ch)
+            images_output.append(fin_img)
+
+        return images_output
+
+def process_image(prompt, image):
+    start_time = time.time()
+    logger.info(f"Processing new request - Prompt: {prompt}, Image size: {image.size if image else 'None'}")
+    
+    try:
+        if image is None:
+            logger.error("No image provided")
+            raise gr.Error("Please upload an image")
+        if not prompt:
+            logger.error("No prompt provided")
+            raise gr.Error("Please enter a prompt")
+            
+        prompt_dict = {'pos': prompt}
+        logger.info("Starting inpainting process")
+        images = inpainter.inpaint(prompt_dict, image)
+        
+        if not images:
+            logger.error("Inpainting failed to produce results")
+            raise gr.Error("Failed to generate images. Please try again.")
+            
+        logger.info(f"Request processed in {time.time() - start_time:.2f} seconds")
+        return images
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise gr.Error(f"Error processing image: {str(e)}")
+
+# Initialize the model
+init()
+
+# Create Gradio interface
+with gr.Blocks(title="Looks.Studio - AI Clothing Inpainting") as demo:
+    gr.Markdown("# Looks.Studio - AI Clothing Inpainting")
+    gr.Markdown("Upload an image and describe the clothing you want to generate")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                type="pil", 
+                label="Upload Image",
+                height=512
+            )
+            prompt = gr.Textbox(label="Describe the clothing you want to generate")
+            generate_btn = gr.Button("Generate")
+        
+        with gr.Column():
+            gallery = gr.Gallery(
+                label="Generated Images", 
+                show_label=False,
+                columns=2,
+                height=512
+            )
+    
+    generate_btn.click(
+        fn=process_image,
+        inputs=[prompt, input_image],
+        outputs=gallery
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True) 

configs/configs.json

@@ -0,0 +1,20 @@
+{
+    "models": {
+        "schp": {
+            "download_url": "https://storage.googleapis.com/platform-ai-prod/looks_studio_data/schp_checkpoint.pth",
+            "path": "checkpoints/schp.pth"
+        },
+        "u2net": {
+            "download_url": "https://storage.googleapis.com/platform-ai-prod/looks_studio_data/cloth_segm_u2net_latest.pth",
+            "path": "checkpoints/cloth_segm_u2net_latest.pth"
+        },
+        "realesrgan": {
+            "download_url": "https://storage.googleapis.com/platform-ai-prod/looks_studio_data/RealESRGAN_x4plus.pth",
+            "path": "checkpoints/realesrgan_x4plus.pth"
+        },
+        "diffuser": {
+            "download_url": "https://storage.googleapis.com/platform-ai-prod/looks_studio_data/diffusers/stable_diffusion_2_checkpoint.tar",
+            "path": "checkpoints/stable_diffusion_2_inpainting"
+        }
+    }
+}

download.py

@@ -0,0 +1,34 @@
+# In this file, we define download_model
+# It runs during container build time to get model weights built into the container
+
+import os
+import wget
+import json
+import tarfile
+import tempfile
+
+def download_models(config):
+    # Download parser checkpoint
+    # wget.download(config['schp']['download_url'],
+    #  os.path.join(os.path.dirname(__file__), config['schp']['path']))
+    wget.download(config['u2net']['download_url'],
+     os.path.join(os.path.dirname(__file__), config['u2net']['path']))
+
+    # Download Super resolution model
+    wget.download(config['realesrgan']['download_url'],
+     os.path.join(os.path.dirname(__file__), config['realesrgan']['path']))
+
+    # Download diffuser model checkpoint
+    _, local_file_name = tempfile.mkstemp()
+    local_file_name += '.tar'
+    wget.download(config['diffuser']['download_url'], local_file_name)
+    tar_file = tarfile.open(local_file_name)
+    tar_file.extractall(os.path.join(os.path.dirname(__file__),'checkpoints/'))
+
+if __name__ == "__main__":
+    config_file = "configs/configs.json"
+    config_file = os.path.join(os.path.dirname(__file__), config_file)
+
+    with open(config_file) as fin:
+        config = json.load(fin)
+    download_models(config['models'])

download_models.py

@@ -0,0 +1,22 @@
+import os
+import time
+import logging
+
+logger = logging.getLogger(__name__)
+
+def download_models():
+    """Download required models for the application"""
+    start_time = time.time()
+    logger.info("Starting model download process")
+    
+    try:
+        # Create models directory if it doesn't exist
+        os.makedirs("models", exist_ok=True)
+        
+        logger.info(f"Model setup completed in {time.time() - start_time:.2f} seconds")
+    except Exception as e:
+        logger.error(f"Error in model setup: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    download_models() 

parser/schp_masker.py

@@ -0,0 +1,200 @@
+import cv2
+import sys
+import numpy as np
+
+import torch
+from torch.utils import data
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+
+from PIL import Image
+from collections import OrderedDict
+
+sys.path.insert(1, './schp')
+from utils.transforms import get_affine_transform
+import networks
+from utils.transforms import transform_logits
+
+class PILImageDataset(data.Dataset):
+    def __init__(self, img_lst=[], input_size=[512, 512], transform=None):
+        self.img_lst = img_lst
+        self.input_size = input_size
+        self.transform = transform
+        self.aspect_ratio = input_size[1] * 1.0 / input_size[0]
+        self.input_size = np.asarray(input_size)
+
+    def __len__(self):
+        return len(self.img_lst)
+
+    def _box2cs(self, box):
+        x, y, w, h = box[:4]
+        return self._xywh2cs(x, y, w, h)
+
+    def _xywh2cs(self, x, y, w, h):
+        center = np.zeros((2), dtype=np.float32)
+        center[0] = x + w * 0.5
+        center[1] = y + h * 0.5
+        if w > self.aspect_ratio * h:
+            h = w * 1.0 / self.aspect_ratio
+        elif w < self.aspect_ratio * h:
+            w = h * self.aspect_ratio
+        scale = np.array([w, h], dtype=np.float32)
+        return center, scale
+
+    def __getitem__(self, index):
+        img = np.array(self.img_lst[index])[:,:,::-1]
+        h, w, _ = img.shape
+
+        # Get person center and scale
+        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
+        r = 0
+        trans = get_affine_transform(person_center, s, r, self.input_size)
+        input = cv2.warpAffine(
+            img,
+            trans,
+            (int(self.input_size[1]), int(self.input_size[0])),
+            flags=cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0, 0, 0))
+
+        input = self.transform(input)
+        meta = {
+            'center': person_center,
+            'height': h,
+            'width': w,
+            'scale': s,
+            'rotation': r
+        }
+
+        return input, meta
+
+PALLETE_DICT = {
+    'Background': [], 
+    'Face': [], 
+    'Upper-clothes':[], 
+    'Dress':[], 
+    'Coat':[], 
+    'Soaks':[],
+    'Pants':[],
+    'Jumpsuits':[],
+    'Scarf':[],
+    'Skirt':[], 
+    'Arm':[],
+    'Leg':[], 
+    'Shoe':[]
+    }
+
+val_list = [[0],[1,4,13],[5],[6],[7],[8],[9],[10],[11],[12],[14,15],[16,17],[18,19]]
+for c,j in enumerate(PALLETE_DICT.keys()):
+  val = val_list[c]
+  pallete = []
+  for i in range(60):
+    if len(val) == 1:
+      if (i >= (val[0]*3)) & (i < ((val[0]+1)*3)):
+        pallete.append(255)
+      else:
+        pallete.append(0)
+    if len(val) == 2:
+      if (i >= (val[0]*3)) & (i < ((val[0]+1)*3)) or (i >= (val[1]*3)) & (i < ((val[1]+1)*3)):
+        pallete.append(255)
+      else:
+        pallete.append(0)
+    if len(val) == 3:
+      if (i >= (val[0]*3)) & (i < ((val[0]+1)*3)) or (i >= (val[1]*3)) & (i < ((val[1]+1)*3)) or (i >= (val[2]*3)) & (i < ((val[2]+1)*3)):
+        pallete.append(255)
+      else:
+        pallete.append(0)
+
+  PALLETE_DICT[j] = pallete
+
+
+DATASET_SETTINGS = {
+    'lip': {
+        'input_size': [473, 473],
+        'num_classes': 20,
+        'label': ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat',
+                  'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm',
+                  'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe']
+    },
+    'atr': {
+        'input_size': [512, 512],
+        'num_classes': 18,
+        'label': ['Background', 'Hat', 'Hair', 'Sunglasses', 'Upper-clothes', 'Skirt', 'Pants', 'Dress', 'Belt',
+                  'Left-shoe', 'Right-shoe', 'Face', 'Left-leg', 'Right-leg', 'Left-arm', 'Right-arm', 'Bag', 'Scarf']
+    },
+    'pascal': {
+        'input_size': [512, 512],
+        'num_classes': 7,
+        'label': ['Background', 'Head', 'Torso', 'Upper Arms', 'Lower Arms', 'Upper Legs', 'Lower Legs'],
+    }
+}
+
+
+
+class SCHPParser:
+    def __init__(self, checkpoint_path, dataset_settings):
+        self.cp_path = checkpoint_path
+        self.ops = []
+        self.num_classes = dataset_settings['lip']['num_classes']
+        self.input_size = dataset_settings['lip']['input_size']
+        self.label = dataset_settings['lip']['label']
+        self.pallete_dict = PALLETE_DICT
+
+
+        self.img_transforms = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
+        ])
+
+        self.model = self.load_model()
+
+
+    def load_model(self):
+        model = networks.init_model('resnet101', num_classes=self.num_classes, pretrained=None)
+        state_dict = torch.load(self.cp_path)['state_dict']
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]  # remove `module.`
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict)
+        model.cuda()
+        model.eval()
+        return model
+    
+    def create_dataloader(self, img_lst):
+        dataset = PILImageDataset(img_lst, input_size=self.input_size, transform=self.img_transforms)
+        # dataset = SimpleFolderDataset('inputs',input_size, transform)
+        dataloader = DataLoader(dataset)
+        return dataloader
+
+    def get_image_masks(self, img_lst): 
+        print("Evaluating total class number {} with {}".format(self.num_classes, self.label))
+        dataloader = self.create_dataloader(img_lst)
+        with torch.no_grad():
+            for batch in dataloader:
+                op_dict = {}
+                image, meta = batch
+                c = meta['center'].numpy()[0]
+                s = meta['scale'].numpy()[0]
+                w = meta['width'].numpy()[0]
+                h = meta['height'].numpy()[0]
+
+                output = self.model(image.cuda())
+                upsample = torch.nn.Upsample(size=self.input_size, mode='bilinear', align_corners=True)
+                upsample_output = upsample(output[0][-1][0].unsqueeze(0))
+                upsample_output = upsample_output.squeeze()
+                upsample_output = upsample_output.permute(1, 2, 0)  # CHW -> HWC
+
+                logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=self.input_size)
+                parsing_result = np.argmax(logits_result, axis=2)
+                output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
+                for loc, key in  enumerate(self.pallete_dict.keys()):
+                    output_img.putpalette(self.pallete_dict[key])
+                    op_dict.update({
+                        key: output_img.convert('L')
+                    })
+                self.ops.append(op_dict)
+        return self.ops
+
+ 
+

parser/segformer_parser.py

@@ -0,0 +1,185 @@
+import torch
+import numpy as np
+from PIL import Image
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+import torch.nn.functional as F
+import logging
+import time
+from typing import Tuple, Optional
+
+logger = logging.getLogger('looks.studio.segformer')
+
+class SegformerParser:
+    def __init__(self, model_path="mattmdjaga/segformer_b2_clothes"):
+        self.start_time = time.time()
+        logger.info(f"Initializing SegformerParser with model: {model_path}")
+        
+        try:
+            self.processor = SegformerImageProcessor.from_pretrained(model_path)
+            self.model = AutoModelForSemanticSegmentation.from_pretrained(model_path)
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            logger.info(f"Using device: {self.device}")
+            self.model.to(self.device)
+            
+            # Define clothing-related labels
+            self.clothing_labels = {
+                4: "upper-clothes",
+                5: "skirt",
+                6: "pants",
+                7: "dress",
+                8: "belt",
+                9: "left-shoe",
+                10: "right-shoe",
+                14: "left-arm",
+                15: "right-arm",
+                17: "scarf"
+            }
+            
+            logger.info(f"SegformerParser initialized in {time.time() - self.start_time:.2f} seconds")
+        except Exception as e:
+            logger.error(f"Failed to initialize SegformerParser: {str(e)}")
+            raise
+    
+    def _resize_image(self, image: Image.Image, max_size: int = 1024) -> Tuple[Image.Image, float]:
+        """Resize image while maintaining aspect ratio if it exceeds max_size"""
+        width, height = image.size
+        scale = 1.0
+        
+        if width > max_size or height > max_size:
+            scale = max_size / max(width, height)
+            new_width = int(width * scale)
+            new_height = int(height * scale)
+            logger.info(f"Resizing image from {width}x{height} to {new_width}x{new_height}")
+            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        
+        return image, scale
+    
+    def _validate_image(self, image: Image.Image) -> bool:
+        """Validate input image"""
+        if not isinstance(image, Image.Image):
+            logger.error("Input is not a PIL Image")
+            return False
+        
+        if image.mode not in ['RGB', 'RGBA']:
+            logger.error(f"Unsupported image mode: {image.mode}")
+            return False
+        
+        width, height = image.size
+        if width < 64 or height < 64:
+            logger.error(f"Image too small: {width}x{height}")
+            return False
+        
+        if width > 4096 or height > 4096:
+            logger.error(f"Image too large: {width}x{height}")
+            return False
+        
+        return True
+    
+    def get_image_mask(self, image: Image.Image) -> Optional[Image.Image]:
+        """Generate segmentation mask for clothing"""
+        start_time = time.time()
+        logger.info(f"Starting segmentation for image size: {image.size}")
+        
+        try:
+            # Validate input image
+            if not self._validate_image(image):
+                return None
+            
+            # Convert RGBA to RGB if necessary
+            if image.mode == 'RGBA':
+                logger.info("Converting RGBA to RGB")
+                image = image.convert('RGB')
+            
+            # Resize image if too large
+            image, scale = self._resize_image(image)
+            
+            # Process the image
+            logger.info("Processing image with Segformer")
+            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+            
+            # Get predictions
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                logits = outputs.logits.cpu()
+                
+                # Upsample logits to original image size
+                upsampled_logits = F.interpolate(
+                    logits,
+                    size=image.size[::-1],
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                
+                # Get the segmentation mask
+                pred_seg = upsampled_logits.argmax(dim=1)[0]
+                
+                # Create a binary mask for clothing
+                mask = torch.zeros_like(pred_seg)
+                for label_id in self.clothing_labels.keys():
+                    mask[pred_seg == label_id] = 255
+                
+                # Convert to PIL Image
+                mask = Image.fromarray(mask.numpy().astype(np.uint8))
+                
+                # Resize mask back to original size if needed
+                if scale != 1.0:
+                    original_size = (int(image.size[0] / scale), int(image.size[1] / scale))
+                    logger.info(f"Resizing mask back to original size: {original_size}")
+                    mask = mask.resize(original_size, Image.Resampling.NEAREST)
+                
+                logger.info(f"Segmentation completed in {time.time() - start_time:.2f} seconds")
+                return mask
+                
+        except Exception as e:
+            logger.error(f"Error during segmentation: {str(e)}")
+            return None 
+
+    def get_all_masks(self, image: Image.Image) -> dict:
+        """Return a dict of binary masks for each clothing part label."""
+        start_time = time.time()
+        logger.info(f"Starting per-part segmentation for image size: {image.size}")
+        masks = {}
+        try:
+            # Validate input image
+            if not self._validate_image(image):
+                return masks
+
+            # Convert RGBA to RGB if necessary
+            if image.mode == 'RGBA':
+                logger.info("Converting RGBA to RGB")
+                image = image.convert('RGB')
+
+            # Resize image if too large
+            image, scale = self._resize_image(image)
+
+            # Process the image
+            logger.info("Processing image with Segformer for all masks")
+            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+
+            # Get predictions
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                logits = outputs.logits.cpu()
+                upsampled_logits = F.interpolate(
+                    logits,
+                    size=image.size[::-1],
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                pred_seg = upsampled_logits.argmax(dim=1)[0]
+
+                # For each clothing label, create a binary mask
+                for label_id, part_name in self.clothing_labels.items():
+                    mask = (pred_seg == label_id).numpy().astype(np.uint8) * 255
+                    mask_img = Image.fromarray(mask)
+                    # Resize mask back to original size if needed
+                    if scale != 1.0:
+                        original_size = (int(image.size[0] / scale), int(image.size[1] / scale))
+                        mask_img = mask_img.resize(original_size, Image.Resampling.NEAREST)
+                    masks[part_name] = mask_img
+
+            logger.info(f"Per-part segmentation completed in {time.time() - start_time:.2f} seconds")
+            return masks
+        except Exception as e:
+            logger.error(f"Error during per-part segmentation: {str(e)}")
+            return masks 

parser/u2net_parser.py

@@ -0,0 +1,55 @@
+import os
+# from tqdm import tqdm
+from PIL import Image
+import numpy as np
+import sys
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from .u2net_cloth_seg.data.base_dataset import Normalize_image
+from .u2net_cloth_seg.utils.saving_utils import load_checkpoint_mgpu
+
+from .u2net_cloth_seg.networks import U2NET
+
+class U2NETParser:
+    def __init__(self, checkpoint_path):
+        self.cp_path = checkpoint_path
+        self.img_transforms = transforms.Compose([
+            transforms.ToTensor(),
+            Normalize_image(0.5, 0.5)
+        ])
+        self.model = self.load_model()
+
+
+    def load_model(self):
+        model = U2NET(in_ch=3, out_ch=4)
+        model = load_checkpoint_mgpu(model, self.cp_path)
+        model = model.to("cuda")
+        model = model.eval()
+        return model
+
+    def get_image_mask(self, img): 
+        # print("Evaluating total class number {} with {}".format(self.num_classes, self.label))
+        img_size = img.size
+        img = img.resize((768, 768), Image.BICUBIC)
+        image_tensor = self.img_transforms(img)
+        image_tensor = torch.unsqueeze(image_tensor, 0)
+
+        with torch.no_grad():
+            output_tensor = self.model(image_tensor.to("cuda"))
+
+        output_tensor = F.log_softmax(output_tensor[0], dim=1)
+        output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
+        output_tensor = torch.squeeze(output_tensor, dim=0)
+        output_tensor = torch.squeeze(output_tensor, dim=0)
+        output_arr = output_tensor.cpu().numpy()
+
+        output_arr[output_arr != 1] = 0
+        output_arr[output_arr == 1] = 255
+
+        output_img = Image.fromarray(output_arr.astype('uint8'), mode='L')
+        output_img = output_img.resize(img_size, Image.BICUBIC)
+
+        return output_img

requirements.txt

@@ -0,0 +1,21 @@
+sanic>=25.3.0
+git+https://github.com/huggingface/diffusers.git#egg=diffusers
+transformers>=4.30.0
+scipy>=1.11.0
+opencv-python>=4.8.0
+wget
+# ninja
+accelerate>=0.24.0
+basicsr>=1.4.2
+ftfy>=6.1.1
+# bitsandbytes
+gradio>=3.50.0
+# natsort
+# https://github.com/metrolobo/xformers_wheels/releases/download/1d31a3ac_various_6/xformers-0.0.14.dev0-cp37-cp37m-linux_x86_64.whl
+torch>=2.0.0
+diffusers>=0.19.0
+Pillow>=9.0.0
+requests>=2.28.0
+numpy>=1.24.0
+huggingface_hub>=0.16.0
+matplotlib>=3.7.0  # For visualization

server.py

@@ -0,0 +1,42 @@
+# Do not edit if deploying to Banana Serverless
+# This file is boilerplate for the http server, and follows a strict interface.
+
+# Instead, edit the init() and inference() functions in app.py
+
+from sanic import Sanic, response
+import subprocess
+import app as user_src
+
+# We do the model load-to-GPU step on server startup
+# so the model object is available globally for reuse
+user_src.init()
+
+# Create the http server app
+server = Sanic("my_app")
+
+# Healthchecks verify that the environment is correct on Banana Serverless
+@server.route('/healthcheck', methods=["GET"])
+def healthcheck(request):
+    # dependency free way to check if GPU is visible
+    gpu = False
+    out = subprocess.run("nvidia-smi", shell=True)
+    if out.returncode == 0: # success state on shell command
+        gpu = True
+
+    return response.json({"state": "healthy", "gpu": gpu})
+
+# Inference POST handler at '/' is called for every http call from Banana
+@server.route('/', methods=["POST"]) 
+def inference(request):
+    try:
+        model_inputs = response.json.loads(request.json)
+    except:
+        model_inputs = request.json
+
+    output = user_src.inference(model_inputs)
+
+    return response.json(output)
+
+
+if __name__ == '__main__':
+    server.run(host='0.0.0.0', port=8000, workers=1)

upscaler/realesrgan_upscaler.py

@@ -0,0 +1,30 @@
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from .real_esrgan.realesrgan import RealESRGANer
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+class RealESRGAN:
+    def __init__(self, checkpoint_path):
+
+        self.netscale = 4
+
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+
+        self.upsampler = RealESRGANer(
+            scale=self.netscale,
+            model_path=checkpoint_path,
+            dni_weight=None,
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=True)
+    
+    def upscale(self, pil_image, scale_factor=3):
+        cv2_img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+        op, _ = self.upsampler.enhance(cv2_img, outscale=scale_factor)
+        pil_image_fin = Image.fromarray(cv2.cvtColor(op, cv2.COLOR_BGR2RGB))
+        return pil_image_fin