Update app.py

app.py

@@ -26,6 +26,18 @@ from PIL import Image
 from gradio_client import Client, handle_file
 import uuid
 
+import gradio as gr
+import spaces
+import torch
+from diffusers import AutoencoderKL, TCDScheduler
+from diffusers.models.model_loading_utils import load_state_dict
+from gradio_imageslider import ImageSlider
+from huggingface_hub import hf_hub_download
+from transformers import pipeline
+
+from controlnet_union import ControlNetModel_Union
+from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
+
 
 def clear_memory():
     """메모리 정리 함수"""
@@ -119,8 +131,94 @@ try:
 except Exception as e:
     print(f"Warning: Could not move pipeline to CUDA: {str(e)}")
 
+
+#------------------------------- 이미지 인페인팅 ----------------------
+
 client = Client("NabeelShar/BiRefNet_for_text_writing")
 
+MODELS = {
+    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
+}
+
+config_file = hf_hub_download(
+    "xinsir/controlnet-union-sdxl-1.0",
+    filename="config_promax.json",
+)
+
+config = ControlNetModel_Union.load_config(config_file)
+controlnet_model = ControlNetModel_Union.from_config(config)
+model_file = hf_hub_download(
+    "xinsir/controlnet-union-sdxl-1.0",
+    filename="diffusion_pytorch_model_promax.safetensors",
+)
+state_dict = load_state_dict(model_file)
+model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
+    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
+)
+model.to(device="cuda", dtype=torch.float16)
+
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
+).to("cuda")
+
+pipe = StableDiffusionXLFillPipeline.from_pretrained(
+    "SG161222/RealVisXL_V5.0_Lightning",
+    torch_dtype=torch.float16,
+    vae=vae,
+    controlnet=model,
+    variant="fp16",
+).to("cuda")
+
+pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+
+def translate_if_korean(text):
+    # 입력된 텍스트가 한글을 포함하고 있는지 확인
+    if any('\u3131' <= char <= '\u318E' or '\uAC00' <= char <= '\uD7A3' for char in text):
+        # 한글이 포함되어 있다면 번역
+        translated = translator(text)[0]['translation_text']
+        print(f"Translated prompt: {translated}")  # 디버깅을 위한 출력
+        return translated
+    return text
+
+@spaces.GPU
+def fill_image(prompt, image, model_selection):
+    # 프롬프트 번역
+    translated_prompt = translate_if_korean(prompt)
+    
+    (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = pipe.encode_prompt(translated_prompt, "cuda", True)
+
+    source = image["background"]
+    mask = image["layers"][0]
+
+    alpha_channel = mask.split()[3]
+    binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
+    cnet_image = source.copy()
+    cnet_image.paste(0, (0, 0), binary_mask)
+
+    for image in pipe(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+        image=cnet_image,
+    ):
+        yield image, cnet_image
+
+    image = image.convert("RGBA")
+    cnet_image.paste(image, (0, 0), binary_mask)
+
+    yield source, cnet_image
+
+def clear_result():
+    return gr.update(value=None)
+
+#--------------- 이미지 인페인팅 끝 ----------------    
+
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name