:tada: change example

app.py

@@ -1,5 +1,3 @@
-import dataclasses
-
 import gradio as gr
 import spaces
 import torch
@@ -7,6 +5,8 @@ from PIL import Image
 from diffusers import DiffusionPipeline
 from diffusers.utils import make_image_grid
 
+from src.example import EXAMPLES
+
 DIFFUSERS_MODEL_IDS = [
     # SD Models
     "stabilityai/stable-diffusion-3-medium-diffusers",
@@ -22,50 +22,36 @@ EXTERNAL_MODEL_MAPPING = {
 }
 MODEL_CHOICES = DIFFUSERS_MODEL_IDS + list(EXTERNAL_MODEL_MAPPING.keys())
 
-current_model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = None
-
-
-@dataclasses.dataclass
-class Input:
-    prompt: str
-    model_id: str = "stabilityai/stable-diffusion-3-medium-diffusers"
-    negative_prompt: str = ''
-    width: int = 1024
-    height: int = 1024
-    guidance_scale: float = 7.5
-    num_inference_step: int = 28
-    num_images: int = 4
-    use_safety_checker: bool = True
-    use_model_offload: bool = False
-    seed: int = 8888
-
-    def to_list(self):
-        return [
-            self.prompt, self.model_id, self.negative_prompt,
-            self.width, self.height, self.guidance_scale,
-            self.num_inference_step, self.num_images, self.use_safety_checker, self.use_model_offload,
-            self.seed
-        ]
 
 
-EXAMPLES = [
-    Input(prompt='A cat holding a sign that says Hello world').to_list(),
-    Input(
-        prompt='Beautiful pixel art of a Wizard with hovering text "Achivement unlocked: Diffusion models can spell now"'
-    ).to_list(),
-    Input(prompt='A corgi wearing sunglasses says "U-Net is OVER!!"').to_list(),
-    Input(
-        prompt='Cinematic Photo of a beautiful korean fashion model bokeh train',
-        model_id='Beautiful Realistic Asians',
-        negative_prompt='worst_quality, BadNegAnatomyV1-neg, bradhands cartoon, cgi, render, illustration, painting, drawing',
-        width=512,
-        height=512,
-        guidance_scale=5.0,
-        num_inference_step=50,
-    ).to_list()
-]
+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='DeepNegativeV1')
+        pipe.load_textual_inversion('checkpoints/embeddings/easynegative.safetensors', token='EasyNegative')
+
+    # VRAM が少ないとき用の対策
+    if use_model_offload:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe = pipe.to(device)
+
+    return pipe
 
 
 @spaces.GPU(duration=120)
@@ -84,37 +70,8 @@ def inference(
         seed: int = 8888,
         progress=gr.Progress(track_tqdm=True),
 ) -> Image.Image:
-    progress(0, "Starting inference...")
-
-    global current_model_id, pipe
-
-    progress(0.1, 'Loading pipeline...')
-    if model_id not in DIFFUSERS_MODEL_IDS:
-        model_id = EXTERNAL_MODEL_MAPPING[model_id]
-
-    pipe = DiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-    )
-
-    current_model_id = model_id
-
-    if not safety_checker:
-        pipe.safety_checker = None
-
-    if model_id not in DIFFUSERS_MODEL_IDS:
-        progress(0.3, 'Loading Textual Inversion...')
-        # 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='DeepNegativeV1')
-        pipe.load_textual_inversion('checkpoints/embeddings/easynegative.safetensors', token='EasyNegative')
-
-    # Generation
-    progress(0.4, 'Generating images...')
-    if use_model_offload:
-        pipe.enable_model_cpu_offload()
-    else:
-        pipe = pipe.to('cuda')
+    progress(0, 'Loading pipeline...')
+    pipe = load_pipeline(model_id, use_model_offload, safety_checker)
 
     generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -137,10 +94,12 @@ def inference(
     return image
 
 
-if __name__ == "__main__":
+def build_interface():
+    """Build Gradio Interface"""
+
     theme = gr.themes.Default(primary_hue=gr.themes.colors.emerald)
 
-    with gr.Blocks(theme=theme) as demo:
+    with gr.Blocks(theme=theme) as interface:
         gr.Markdown(f"# Stable Diffusion Demo")
 
         with gr.Row():
@@ -217,5 +176,9 @@ if __name__ == "__main__":
             fn=inference,
             cache_examples='lazy'
         )
+    return interface
 
-    demo.queue().launch()
+
+if __name__ == "__main__":
+    iface = build_interface()
+    iface.queue().launch()

src/example.py

@@ -0,0 +1,46 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class Example:
+    prompt: str
+    model_id: str = "stabilityai/stable-diffusion-3-medium-diffusers"
+    negative_prompt: str = ''
+    width: int = 1024
+    height: int = 1024
+    guidance_scale: float = 7.5
+    num_inference_step: int = 28
+    num_images: int = 4
+    use_safety_checker: bool = True
+    use_model_offload: bool = False
+    seed: int = 8888
+
+    def to_list(self):
+        return [
+            self.prompt, self.model_id, self.negative_prompt,
+            self.width, self.height, self.guidance_scale,
+            self.num_inference_step, self.num_images, self.use_safety_checker, self.use_model_offload,
+            self.seed
+        ]
+
+
+EXAMPLES = [
+    Example(
+        prompt='A cat holding a sign that says Hello world'
+    ).to_list(),
+    Example(
+        prompt='Beautiful pixel art of a Wizard with hovering text "Achivement unlocked: Diffusion models can spell now"'
+    ).to_list(),
+    Example(
+        prompt='A corgi wearing sunglasses says "U-Net is OVER!!"'
+    ).to_list(),
+    Example(
+        prompt='Cinematic Photo of a beautiful korean fashion model bokeh train',
+        model_id='Beautiful Realistic Asians',
+        negative_prompt='worst_quality, BadNegAnatomyV1-neg, bradhands cartoon, cgi, render, illustration, painting, drawing',
+        width=512,
+        height=512,
+        guidance_scale=5.0,
+        num_inference_step=50,
+    ).to_list()
+]