ref

README.md

@@ -28,9 +28,9 @@ a handsome man relaxing in a chair, shirt widely unbuttoned, eyes closed, close
 - [x] adaptive resize размеры
 - [x] пофиксить загрузку маски
 - [x] подобрать нормальные промпты к примерам, добавить negative prompt (https://blog.segmind.com/prompt-guide-for-stable-diffusion-xl-crafting-textual-descriptions-for-image-generation/)
-- [ ] разобраться с выставлением device в пайплайн
-- [ ] удалить ненужные веса
-- [ ] небольшой рефач - dark theme - можно попробовать
+- [x] удалить ненужные веса
+- [x] небольшой рефач - dark theme - можно попробовать
 - [ ] проверка ссылок
 - [ ] залить версию в main  
+- [ ] разобраться с выставлением device в пайплайн
 - [ ] настроить запуск в hugging space

app.py

@@ -41,12 +41,10 @@ def process(
         output.generated_image,
         output.control_mask,
     ]
-    # return [input_image, input_image]
 
 
 def read_content(file_path: str) -> str:
-    """read the content of target file
-    """
+    """Read the content of target file."""
     with open(file_path, 'r', encoding='utf-8') as f:
         content = f.read()
 
@@ -63,36 +61,36 @@ examples = [[image, prompt[0], prompt[1]] for image, prompt in zip(image_list, p
 block = gr.Blocks().queue()
 with block:
     with gr.Row():
-        gr.HTML(read_content("header.html"))
+        gr.HTML(read_content('header.html'))
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="numpy")
-            prompt = gr.Textbox(label="Prompt")
-            negative_prompt = gr.Textbox(label="Negative Prompt")
+            input_image = gr.Image(type='numpy')
+            prompt = gr.Textbox(label='Prompt')
+            negative_prompt = gr.Textbox(label='Negative Prompt')
             with gr.Row():
-                generate_from_mask = gr.Checkbox(label="Input image is already a control mask", value=False)
-            run_button = gr.Button(value="Run")
-            with gr.Accordion("Advanced options", open=False):
+                generate_from_mask = gr.Checkbox(label='Input image is already a control mask', value=False)
+            run_button = gr.Button(value='Run')
+            with gr.Accordion('Advanced options', open=False):
                 target_image_size = gr.Slider(
-                    label="Image target size:",
+                    label='Image target size:',
                     minimum=512,
                     maximum=2048,
                     value=768,
                     step=64,
                 )
                 max_image_size = gr.Slider(
-                    label="Image max size:",
+                    label='Image max size:',
                     minimum=512,
                     maximum=2048,
                     value=1024,
                     step=64,
                 )
-                num_inference_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
-                conditioning_scale = gr.Slider(label="Conditioning scale", minimum=0.0, maximum=5.0, value=1.0, step=0.1)
-                seed = gr.Slider(label="Seed", minimum=0, maximum=config.max_seed, step=1, value=0)
+                num_inference_steps = gr.Slider(label='Number of steps', minimum=1, maximum=100, value=20, step=1)
+                guidance_scale = gr.Slider(label='Guidance scale', minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                conditioning_scale = gr.Slider(label='Conditioning scale', minimum=0.0, maximum=5.0, value=1.0, step=0.1)
+                seed = gr.Slider(label='Seed', minimum=0, maximum=config.max_seed, step=1, value=0)
 
-            gr.Examples(examples=examples, inputs=[input_image, prompt, negative_prompt], label="Examples - Input Images", examples_per_page=12)
+            gr.Examples(examples=examples, inputs=[input_image, prompt, negative_prompt], label='Examples - Input Images', examples_per_page=12)
 
             gr.HTML(
                 """
@@ -105,8 +103,8 @@ with block:
                 """)
 
         with gr.Column():
-            generated_output = gr.Image(label="Generated", type="numpy", elem_id="generated")
-            mask_output = gr.Image(label="Mask", type="numpy", elem_id="mask")
+            generated_output = gr.Image(label='Generated', type='numpy', elem_id='generated')
+            mask_output = gr.Image(label='Mask', type='numpy', elem_id='mask')
 
     ips = [input_image, prompt, negative_prompt, generate_from_mask, num_inference_steps, guidance_scale, conditioning_scale, target_image_size, max_image_size, seed]
     run_button.click(fn=process, inputs=ips, outputs=[generated_output, mask_output])

src/inference_no_condition.py

@@ -6,13 +6,13 @@ import matplotlib.pyplot as plt
 
 
 pipe = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
+    'stabilityai/stable-diffusion-xl-base-1.0',
     torch_dtype=torch.float16,
     use_safetensors=True,
-    variant="fp16",
+    variant='fp16',
 )
 
-pipe.to("cuda")
+pipe.to('cuda')
 
 with open('../examples/prompts.json', 'r') as f:
     prompts_list = list(json.load(f).values())

src/pipeline.py

@@ -9,11 +9,9 @@ import torch
 
 from src.preprocess import HWC3
 from src.unet.predictor import generate_mask, load_seg_model
-
 from config import PipelineConfig
 
 
-
 @dataclass
 class PipelineOutput:
     control_mask: np.ndarray
@@ -49,7 +47,7 @@ class FashionPipeline:
         max_image_size: int,
         seed: int,
     ) -> PipelineOutput:
-
+        """Runs image generation pipeline."""
         # check image format
         control_image = HWC3(control_image)
 
@@ -99,7 +97,10 @@ class FashionPipeline:
         max_image_size: int = 768,
         divisible: int = 64,
     ) -> Image:
-
+        """Resizes the image so that width and height are
+        divided by 'divisible' while maintaining aspect ratio.
+        Restrict image size with target_image_size and max_image_size.
+        """
         assert target_image_size % divisible == 0
         assert max_image_size % divisible == 0
         assert max_image_size >= target_image_size
@@ -130,15 +131,13 @@ class FashionPipeline:
         self.controlnet = ControlNetModel.from_pretrained(
             self.config.controlnet_path,
             torch_dtype=torch.float16,
-            # device_map="auto",
-        )
+        ).to(self.device)
 
         self.pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
             self.config.base_model_path,
             controlnet=self.controlnet,
             torch_dtype=torch.float16,
-            # device_map="auto",
-        )
+        ).to(self.device)
 
         self.pipeline.scheduler = UniPCMultistepScheduler.from_config(self.pipeline.scheduler.config)