Add a Dockfile and Add Fast API function in app.py

Dockerfile

@@ -0,0 +1,30 @@
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements file
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the application code
+COPY . .
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1
+ENV CUDA_VISIBLE_DEVICES=0
+
+# Expose port (if your app needs it)
+EXPOSE 7860
+
+# Command to run the application
+CMD ["python", "app.py"] 

app.py

@@ -1,5 +1,10 @@
 import spaces
-import gradio as gr
+from fastapi import FastAPI, File, UploadFile
+import pickle
+import uvicorn
+import pandas as pd
+import io
+# import gradio as gr
 from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
@@ -269,57 +274,83 @@ for ex_human in human_list_path:
 ##default human
 
 
-image_blocks = gr.Blocks().queue()
-with image_blocks as demo:
-    gr.Markdown("👕👔👚 GSR 임성남 👕👔👚")
-    gr.Markdown("생성형 AI를 활용한 가상 의상 착장 Protype (Based on IDM-VTON)!!!")
-    with gr.Row():
-        with gr.Column():
-            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
-            with gr.Row():
-                is_checked = gr.Checkbox(label="'상의' 자동 마스킹", info="체크 '해제'해야 '마스킹 의류 선택'이 반영됩니다.",value=True)
-            # with gr.Row():
-                # gr.Radio(["수정중",], label="(수정중)마스킹 의류 선택", info="자동으로 마스킹할 의류 위치를 선택하세요")
-                # is_category = gr.Radio(["upper_body", "lower_body", "dresses","layer"], label="마스킹 의류 선택", info="자동으로 마스킹할 의류 위치를 선택하세요")
-            with gr.Row():
-                is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
-
-            example = gr.Examples(
-                inputs=imgs,
-                examples_per_page=10,
-                examples=human_ex_list
-            )
-
-        with gr.Column():
-            garm_img = gr.Image(label="Garment", sources='upload', type="pil")
-            with gr.Row(elem_id="prompt-container"):
-                with gr.Row():
-                    prompt = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
-            example = gr.Examples(
-                inputs=garm_img,
-                examples_per_page=8,
-                examples=garm_list_path)
-        with gr.Column():
-            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
-            masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
-        with gr.Column():
-            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
-            image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)
-
-
-
-
-    with gr.Column():
-        try_button = gr.Button(value="Try-on")
-        with gr.Accordion(label="Advanced Settings", open=False):
-            with gr.Row():
-                denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1)
-                seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
-
-
-    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
+# API base configuration
+app = FastAPI()
+
+@app.get("/")
+def api_home():
+    return {"message": "Hello World"}
+
+
+#AI 모델 실행 API, 처리완료된 이미지를 리턴한다. 
+@app.post("/vton/")
+async def vton_run(
+    upload_human: UploadFile = File(...),
+    upload_cloth: UploadFile = File(...),
+    input_prompt: str = None,
+    is_checked: bool = True,
+    is_checked_crop: bool = True,
+    denoise_steps: int = 30,
+    seed: int = 42
+):
+    target_human = Image.open(io.BytesIO(await upload_human.read()))
+    target_cloth = Image.open(io.BytesIO(await upload_cloth.read()))
+
+    results = start_tryon(target_human, target_cloth, input_prompt, is_checked, is_checked_crop, denoise_steps, seed)
+    return results[0]
+
+
+# image_blocks = gr.Blocks().queue()
+# with image_blocks as demo:
+#     gr.Markdown("👕👔👚 GSR 임성남 👕👔👚")
+#     gr.Markdown("생성형 AI를 활용한 가상 의상 착장 Protype (Based on IDM-VTON)!!!")
+#     with gr.Row():
+#         with gr.Column():
+#             imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
+#             with gr.Row():
+#                 is_checked = gr.Checkbox(label="'상의' 자동 마스킹", info="체크 '해제'해야 '마스킹 의류 선택'이 반영됩니다.",value=True)
+#             # with gr.Row():
+#                 # gr.Radio(["수정중",], label="(수정중)마스킹 의류 선택", info="자동으로 마스킹할 의류 위치를 선택하세요")
+#                 # is_category = gr.Radio(["upper_body", "lower_body", "dresses","layer"], label="마스킹 의류 선택", info="자동으로 마스킹할 의류 위치를 선택하세요")
+#             with gr.Row():
+#                 is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
+
+#             example = gr.Examples(
+#                 inputs=imgs,
+#                 examples_per_page=10,
+#                 examples=human_ex_list
+#             )
+
+#         with gr.Column():
+#             garm_img = gr.Image(label="Garment", sources='upload', type="pil")
+#             with gr.Row(elem_id="prompt-container"):
+#                 with gr.Row():
+#                     prompt = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
+#             example = gr.Examples(
+#                 inputs=garm_img,
+#                 examples_per_page=8,
+#                 examples=garm_list_path)
+#         with gr.Column():
+#             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+#             masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
+#         with gr.Column():
+#             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+#             image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)
+
+
+
+
+#     with gr.Column():
+#         try_button = gr.Button(value="Try-on")
+#         with gr.Accordion(label="Advanced Settings", open=False):
+#             with gr.Row():
+#                 denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1)
+#                 seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
+
+
+#     try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
 
             
 
 
-image_blocks.launch()
+# image_blocks.launch()