app.py

app.py

@@ -1,99 +1,98 @@
 import spaces
 import gradio as gr
 import torch
-
-torch.jit.script = lambda f: f
-from diffusers import (
-    ControlNetModel,
-    StableDiffusionXLControlNetImg2ImgPipeline,
-    DDIMScheduler,
-)
+from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, DDIMScheduler
 from compel import Compel, ReturnedEmbeddingsType
 from PIL import Image
 import os
 import time
-import numpy as np
 
 from utils.utils import load_cn_model, load_cn_config, load_tagger_model, load_lora_model, resize_image_aspect_ratio, base_generation
 from utils.prompt_analysis import PromptAnalysis
 
-path = os.getcwd()
-cn_dir = f"{path}/controlnet"
-os.makedirs(cn_dir)
-tagger_dir = f"{path}/tagger"
-os.mkdir(tagger_dir)
-lora_dir = f"{path}/lora"
-os.mkdir(lora_dir)
-
-load_cn_model(cn_dir)
-load_cn_config(cn_dir)
-load_tagger_model(tagger_dir)
-load_lora_model(lora_dir)
-
-IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
-IS_SPACE = os.environ.get("SPACE_ID", None) is not None
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16
-
-LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
-
-print(f"device: {device}")
-print(f"dtype: {dtype}")
-print(f"low memory: {LOW_MEMORY}")
-
-
-model = "cagliostrolab/animagine-xl-3.1"
-scheduler = DDIMScheduler.from_pretrained(model, subfolder="scheduler")
-controlnet = ControlNetModel.from_pretrained(cn_dir, torch_dtype=torch.float16, use_safetensors=True)
-pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-    model,
-    controlnet=controlnet,
-    torch_dtype=dtype,
-    use_safetensors=True,
-    scheduler=scheduler,
-)
-
-pipe.load_lora_weights(
-    lora_dir, 
-    weight_name="sdxl_BWLine.safetensors"
-)
-
-
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True],
-)
-pipe = pipe.to(device)
-
-
-
 class Img2Img:
     def __init__(self):
-        self.input_image_path = None
+        self.setup_paths()
+        self.setup_models()
+        self.compel = self.setup_compel()
+        self.demo = self.layout()
+
+    def setup_paths(self):
+        self.path = os.getcwd()
+        self.cn_dir = f"{self.path}/controlnet"
+        self.tagger_dir = f"{self.path}/tagger"
+        self.lora_dir = f"{self.path}/lora"
+        os.makedirs(self.cn_dir, exist_ok=True)
+        os.makedirs(self.tagger_dir, exist_ok=True)
+        os.makedirs(self.lora_dir, exist_ok=True)
+
+    def setup_models(self):
+        load_cn_model(self.cn_dir)
+        load_cn_config(self.cn_dir)
+        load_tagger_model(self.tagger_dir)
+        load_lora_model(self.lora_dir)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16
+        self.model = "cagliostrolab/animagine-xl-3.1"
+        self.scheduler = DDIMScheduler.from_pretrained(self.model, subfolder="scheduler")
+        self.controlnet = ControlNetModel.from_pretrained(self.cn_dir, torch_dtype=self.dtype, use_safetensors=True)
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            self.model,
+            controlnet=self.controlnet,
+            torch_dtype=self.dtype,
+            use_safetensors=True,
+            scheduler=self.scheduler,
+        )
+        self.pipe.load_lora_weights(self.lora_dir, weight_name="sdxl_BWLine.safetensors")
+        self.pipe = self.pipe.to(self.device)
+
+    def setup_compel(self):
+        return Compel(
+            tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2],
+            text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2],
+            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+            requires_pooled=[False, True],
+        )
+
+    def layout(self):
+        css = """
+        #intro{
+            max-width: 32rem;
+            text-align: center;
+            margin: 0 auto;
+        }
+        """
+        with gr.Blocks(css=css) as demo:
+            with gr.Row():
+                with gr.Column():
+                    self.input_image_path = gr.Image(label="入力画像", type='filepath')
+                    self.prompt_analysis = PromptAnalysis(self.tagger_dir)
+                    self.prompt, self.negative_prompt = self.prompt_analysis.layout(self.input_image_path)
+                    self.controlnet_scale = gr.Slider(minimum=0.5, maximum=1.25, value=1.0, step=0.01, label="線画忠実度")
+                    generate_button = gr.Button("生成")
+                with gr.Column():
+                    self.output_image = gr.Image(type="pil", label="生成画像")
+
+            generate_button.click(
+                fn=self.predict,
+                inputs=[self.input_image_path, self.prompt, self.negative_prompt, self.controlnet_scale],
+                outputs=self.output_image
+            )
+        return demo
 
     @spaces.GPU
-    def predict(
-        self,
-        input_image_path,
-        prompt,
-        negative_prompt,
-        controlnet_conditioning_scale,
-    ):
+    def predict(self, input_image_path, prompt, negative_prompt, controlnet_scale):
         input_image_pil = Image.open(input_image_path)
-        base_size =input_image_pil.size
-        resize_image= resize_image_aspect_ratio(input_image_pil)
+        base_size = input_image_pil.size
+        resize_image = resize_image_aspect_ratio(input_image_pil)
         resize_image_size = resize_image.size
-        width = resize_image_size[0]
-        height = resize_image_size[1]
+        width, height = resize_image_size
         white_base_pil = base_generation(resize_image.size, (255, 255, 255, 255)).convert("RGB")
-        conditioning, pooled = compel([prompt, negative_prompt])
+        conditioning, pooled = self.compel([prompt, negative_prompt])
         generator = torch.manual_seed(0)
         last_time = time.time()
 
-        output_image = pipe(
+        output_image = self.pipe(
             image=white_base_pil,
             control_image=resize_image,
             strength=1.0,
@@ -103,7 +102,7 @@ class Img2Img:
             negative_pooled_prompt_embeds=pooled[1:2],
             width=width,
             height=height,
-            controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+            controlnet_conditioning_scale=float(controlnet_scale),
             controlnet_start=0.0,
             controlnet_end=1.0,
             generator=generator,
@@ -115,44 +114,5 @@ class Img2Img:
         output_image = output_image.resize(base_size, Image.LANCZOS)
         return output_image
 
-
-    css = """
-    #intro{
-        # max-width: 32rem;
-        # text-align: center;
-        # margin: 0 auto;
-    }
-    """
-    def layout(self,css):
-        with gr.Blocks(css=css) as demo:
-            with gr.Column():
-                # 画像アップロード用の行
-                with gr.Row():
-                    with gr.Column():
-                        self.input_image_path = gr.Image(label="入力画像",  type='filepath')
-                
-                # プロンプト入力用の行
-                with gr.Row():
-                    prompt_analysis = PromptAnalysis(tagger_dir)
-                    [prompt, nega] = prompt_analysis.layout(self.input_image_path)           
-                # 画像の詳細設定用のスライダー行
-                with gr.Row():
-                    controlnet_conditioning_scale = gr.Slider(minimum=0.5, maximum=1.25, value=1.0, step=0.01, interactive=True, label="線画忠実度")
-            
-                # 画像生成ボタンの行
-                with gr.Row():
-                    generate_button = gr.Button("生成", interactive=False)
-
-            with gr.Column():
-                output_image = gr.Image(type="pil", label="Output Image")
-          
-            # ボタンのクリックイベントを設定
-            generate_button.click(
-                fn=self.predict,
-                inputs=[self.input_image_path, prompt, nega, controlnet_conditioning_scale],
-                outputs=[output_image]
-            )
-
-        # デモの設定と起動
-        demo.queue(api_open=True)
-        demo.launch(show_api=True)
+img2img = Img2Img()
+img2img.demo.launch()