:art: Improve structure

app.py

@@ -1,125 +1,8 @@
 import gradio as gr
-import spaces
-import torch
-from PIL import Image
-from compel import Compel, DiffusersTextualInversionManager
-from diffusers import DiffusionPipeline
-from diffusers.utils import make_image_grid
 
+from src.const import MODEL_CHOICES
 from src.example import EXAMPLES
-
-DIFFUSERS_MODEL_IDS = [
-    # SD Models
-    "stabilityai/stable-diffusion-3-medium-diffusers",
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    "stabilityai/stable-diffusion-2-1",
-    "runwayml/stable-diffusion-v1-5",
-
-    # Other Models
-    "Prgckwb/trpfrog-diffusion",
-]
-EXTERNAL_MODEL_MAPPING = {
-    "Beautiful Realistic Asians": "checkpoints/diffusers/Beautiful Realistic Asians v7",
-}
-MODEL_CHOICES = DIFFUSERS_MODEL_IDS + list(EXTERNAL_MODEL_MAPPING.keys())
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-def load_pipeline(model_id, use_model_offload, safety_checker):
-    # Diffusers リポジトリ内のモデル
-    if model_id in DIFFUSERS_MODEL_IDS:
-        pipe = DiffusionPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-        )
-
-    # CIVITAI 系列由来のモデル
-    else:
-        pipe = DiffusionPipeline.from_pretrained(
-            EXTERNAL_MODEL_MAPPING[model_id],
-            torch_dtype=torch.float16,
-        )
-
-        # Load Textual Inversion
-        pipe.load_textual_inversion("checkpoints/embeddings/BadNegAnatomyV1 neg.pt", token='BadNegAnatomyV1-neg')
-        pipe.load_textual_inversion("checkpoints/embeddings/Deep Negative V1 75T.pt", token='DeepNegative')
-        pipe.load_textual_inversion("checkpoints/embeddings/easynegative.safetensors", token='EasyNegative')
-        pipe.load_textual_inversion("checkpoints/embeddings/Negative Hand Embedding.pt", token='negative_hand-neg')
-
-        # Load LoRA
-        pipe.load_lora_weights("checkpoints/lora/detailed style SD1.5.safetensors", adapter_name='detail')
-        pipe.load_lora_weights("checkpoints/lora/perfection style SD1.5.safetensors", adapter_name='perfection')
-        pipe.load_lora_weights("checkpoints/lora/Hand v3 SD1.5.safetensors", adapter_name='hands')
-        pipe.set_adapters(['detail', 'hands'], adapter_weights=[0.5, 0.5])
-
-    # VRAM が少ないとき用の対策
-    if use_model_offload:
-        pipe.enable_model_cpu_offload()
-    else:
-        pipe = pipe.to(device)
-
-    if not safety_checker:
-        pipe.safety_checker = None
-
-    return pipe
-
-
-@spaces.GPU(duration=120)
-@torch.inference_mode()
-def inference(
-        prompt: str,
-        model_id: str = "stabilityai/stable-diffusion-3-medium-diffusers",
-        negative_prompt: str = "",
-        width: int = 512,
-        height: int = 512,
-        guidance_scale: float = 7.5,
-        num_inference_steps: int = 50,
-        num_images: int = 4,
-        safety_checker: bool = True,
-        use_model_offload: bool = False,
-        seed: int = 8888,
-        progress=gr.Progress(track_tqdm=True),
-) -> Image.Image:
-    progress(0, 'Loading pipeline...')
-    pipe = load_pipeline(model_id, use_model_offload, safety_checker)
-
-    # For Compel
-    textual_inversion_manager = DiffusersTextualInversionManager(pipe)
-    compel_procs = Compel(
-        tokenizer=pipe.tokenizer,
-        text_encoder=pipe.text_encoder,
-        textual_inversion_manager=textual_inversion_manager,
-        truncate_long_prompts=False,
-    )
-    prompt_embed = compel_procs(prompt)
-    negative_prompt_embed = compel_procs(negative_prompt)
-
-    prompt_embed, negative_prompt_embed = compel_procs.pad_conditioning_tensors_to_same_length(
-        [prompt_embed, negative_prompt_embed]
-    )
-
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    progress(0.3, 'Generating images...')
-    images = pipe(
-        prompt_embeds=prompt_embed,
-        negative_prompt_embeds=negative_prompt_embed,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        num_images_per_prompt=num_images,
-        generator=generator,
-    ).images
-
-    progress(0.9, f'Done generating {num_images} images')
-    if num_images % 2 == 1:
-        image = make_image_grid(images, rows=num_images, cols=1)
-    else:
-        image = make_image_grid(images, rows=2, cols=num_images // 2)
-
-    return image
+from src.inference import inference
 
 
 def build_interface():

src/const.py

@@ -0,0 +1,18 @@
+import torch
+
+DIFFUSERS_MODEL_IDS = [
+    # SD Models
+    "stabilityai/stable-diffusion-3-medium-diffusers",
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "stabilityai/stable-diffusion-2-1",
+    "runwayml/stable-diffusion-v1-5",
+
+    # Other Models
+    "Prgckwb/trpfrog-diffusion",
+]
+EXTERNAL_MODEL_MAPPING = {
+    "Beautiful Realistic Asians": "checkpoints/diffusers/Beautiful Realistic Asians v7",
+}
+MODEL_CHOICES = DIFFUSERS_MODEL_IDS + list(EXTERNAL_MODEL_MAPPING.keys())
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

src/inference.py

@@ -0,0 +1,119 @@
+import gradio as gr
+import spaces
+import torch
+from PIL import Image
+from compel import Compel, DiffusersTextualInversionManager
+from diffusers import DiffusionPipeline, StableDiffusionPipeline
+from diffusers.utils import make_image_grid
+
+from src.const import DIFFUSERS_MODEL_IDS, EXTERNAL_MODEL_MAPPING, DEVICE
+
+
+def load_pipeline(model_id, use_model_offload, safety_checker):
+    # Diffusers リポジトリ内のモデル
+    if model_id in DIFFUSERS_MODEL_IDS:
+        pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+        )
+
+    # CIVITAI 系列由来のモデル
+    else:
+        pipe = DiffusionPipeline.from_pretrained(
+            EXTERNAL_MODEL_MAPPING[model_id],
+            torch_dtype=torch.float16,
+        )
+
+        # Load Textual Inversion
+        pipe.load_textual_inversion("checkpoints/embeddings/BadNegAnatomyV1 neg.pt", token='BadNegAnatomyV1-neg')
+        pipe.load_textual_inversion("checkpoints/embeddings/Deep Negative V1 75T.pt", token='DeepNegative')
+        pipe.load_textual_inversion("checkpoints/embeddings/easynegative.safetensors", token='EasyNegative')
+        pipe.load_textual_inversion("checkpoints/embeddings/Negative Hand Embedding.pt", token='negative_hand-neg')
+
+        # Load LoRA
+        pipe.load_lora_weights("checkpoints/lora/detailed style SD1.5.safetensors", adapter_name='detail')
+        pipe.load_lora_weights("checkpoints/lora/perfection style SD1.5.safetensors", adapter_name='perfection')
+        pipe.load_lora_weights("checkpoints/lora/Hand v3 SD1.5.safetensors", adapter_name='hands')
+        pipe.set_adapters(['detail', 'hands'], adapter_weights=[0.5, 0.5])
+
+    # VRAM が少ないとき用の対策
+    if use_model_offload:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe = pipe.to(DEVICE)
+
+    if not safety_checker:
+        pipe.safety_checker = None
+
+    return pipe
+
+
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def inference(
+        prompt: str,
+        model_id: str = "stabilityai/stable-diffusion-3-medium-diffusers",
+        negative_prompt: str = "",
+        width: int = 512,
+        height: int = 512,
+        guidance_scale: float = 7.5,
+        num_inference_steps: int = 50,
+        num_images: int = 4,
+        safety_checker: bool = True,
+        use_model_offload: bool = False,
+        seed: int = 8888,
+        progress=gr.Progress(track_tqdm=True),
+) -> Image.Image:
+    progress(0, 'Loading pipeline...')
+    pipe = load_pipeline(model_id, use_model_offload, safety_checker)
+
+    # Seed 固定
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+
+    if isinstance(pipe, StableDiffusionPipeline):
+        # For Compel
+        textual_inversion_manager = DiffusersTextualInversionManager(pipe)
+        compel_procs = Compel(
+            tokenizer=pipe.tokenizer,
+            text_encoder=pipe.text_encoder,
+            textual_inversion_manager=textual_inversion_manager,
+            truncate_long_prompts=False,
+        )
+        prompt_embed = compel_procs(prompt)
+        negative_prompt_embed = compel_procs(negative_prompt)
+
+        prompt_embed, negative_prompt_embed = compel_procs.pad_conditioning_tensors_to_same_length(
+            [prompt_embed, negative_prompt_embed]
+        )
+
+        progress(0.3, 'Generating images...')
+        images = pipe(
+            prompt_embeds=prompt_embed,
+            negative_prompt_embeds=negative_prompt_embed,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            num_images_per_prompt=num_images,
+            generator=generator,
+        ).images
+    else:
+        progress(0.3, 'Generating images...')
+        images = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            num_images=num_images,
+            generator=generator,
+        ).images
+
+    progress(0.9, f'Done generating {num_images} images')
+    if num_images % 2 == 1:
+        image = make_image_grid(images, rows=num_images, cols=1)
+    else:
+        image = make_image_grid(images, rows=2, cols=num_images // 2)
+
+    return image