Upload 7 files

app.py

@@ -72,4 +72,4 @@ with gr.Blocks(css='style.css') as demo:
             ''')
             create_upload_demo(HF_TOKEN)
 
-demo.queue(max_size=1).launch(share=False)
+demo.queue(max_size=1).launch(share=True)

app_inference.py

@@ -91,9 +91,13 @@ def create_inference_demo(pipe: InferencePipeline,
                             open=False):
                         with gr.Row():
                             target_image = gr.Image(
-                                label='Target Image', interactive=False)
+                                label='Target Image', 
+                                interactive=False,
+                                type='filepath')
                             target_mask = gr.Image(
-                                label='Target Mask', interactive=False)
+                                label='Target Mask', 
+                                interactive=False,
+                                type='filepath')
                 seed = gr.Slider(label='Seed',
                                  minimum=0,
                                  maximum=100000,

inference.py

@@ -4,8 +4,8 @@ import gc
 import pathlib
 
 import gradio as gr
-import PIL.Image
 import torch
+from PIL import Image, ImageFilter
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 from huggingface_hub import ModelCard
 
@@ -60,22 +60,32 @@ class InferencePipeline:
     def run(
         self,
         model_id: str,
+        seed: int,
         target_image: str,
         target_mask: str,
-        seed: int,
         n_steps: int,
         guidance_scale: float,
-    ) -> PIL.Image.Image:
+    ) -> Image.Image:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
 
         self.load_pipe(model_id)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        image, mask_image = Image.open(target_image), Image.open(target_mask)
+        image, mask_image = image.convert("RGB"), mask_image.convert("L")
+
+        erode_kernel = ImageFilter.MaxFilter(3)
+        mask_image = mask_image.filter(erode_kernel)
+             
+        blur_kernel = ImageFilter.BoxBlur(1)
+        mask_image = mask_image.filter(blur_kernel)
+        
         out = self.pipe(
             "a photo of sks",
-            image=target_image,
-            mask_image=target_mask,
+            image=image,
+            mask_image=mask_image,
             num_inference_steps=n_steps,
             guidance_scale=guidance_scale,
             generator=generator,

train_realfill.py

@@ -16,7 +16,6 @@ import transformers
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
-from huggingface_hub import create_repo, upload_folder
 from packaging import version
 from PIL import Image
 from PIL.ImageOps import exif_transpose
@@ -60,36 +59,6 @@ def make_mask(images, resolution, times=30):
     mask = 1 - mask if random.random() < 0.5 else mask
     return mask
 
-def save_model_card(
-    repo_id: str,
-    base_model: str,
-    target_image: str,
-    target_mask: str,
-    repo_folder=None,
-):
-    yaml = f"""
----
-license: creativeml-openrail-m
-base_model: {base_model}
-target_image: {target_image}
-target_mask: {target_mask}
-tags:
-- stable-diffusion-inpainting
-- stable-diffusion-inpainting-diffusers
-- text-to-image
-- diffusers
-- realfill
-inference: true
----
-    """
-    model_card = f"""
-# RealFill - {repo_id}
-
-This is a realfill model derived from {base_model}. The weights were trained using [RealFill](https://realfill.github.io/).
-"""
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
-
 def log_validation(
     text_encoder,
     tokenizer,
@@ -306,14 +275,6 @@ def parse_args(input_args=None):
     parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
     parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
     parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
-    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
-    parser.add_argument(
-        "--hub_model_id",
-        type=str,
-        default=None,
-        help="The name of the repository to keep in sync with the local `output_dir`.",
-    )
     parser.add_argument(
         "--logging_dir",
         type=str,
@@ -559,11 +520,6 @@ def main(args):
         if args.output_dir is not None:
             os.makedirs(args.output_dir, exist_ok=True)
 
-        if args.push_to_hub:
-            repo_id = create_repo(
-                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
-            ).repo_id
-
     # Load the tokenizer
     if args.tokenizer_name:
         tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
@@ -588,7 +544,7 @@ def main(args):
     config = LoraConfig(
         r=args.lora_rank,
         lora_alpha=args.lora_alpha,
-        target_modules=["to_k", "to_q", "to_v", "key", "query", "value"],
+        target_modules=["to_k", "to_q", "to_v", "out.0"],
         lora_dropout=args.lora_dropout,
         bias=args.lora_bias,
     )
@@ -597,7 +553,7 @@ def main(args):
     config = LoraConfig(
         r=args.lora_rank,
         lora_alpha=args.lora_alpha,
-        target_modules=["k_proj", "q_proj", "v_proj"],
+        target_modules=["k_proj", "q_proj", "v_proj", "out_proj"],
         lora_dropout=args.lora_dropout,
         bias=args.lora_bias,
     )
@@ -930,20 +886,6 @@ def main(args):
             global_step,
         )
 
-        if args.push_to_hub:
-            save_model_card(
-                repo_id,
-                images=images,
-                base_model=args.pretrained_model_name_or_path,
-                repo_folder=args.output_dir,
-            )
-            upload_folder(
-                repo_id=repo_id,
-                folder_path=args.output_dir,
-                commit_message="End of training",
-                ignore_patterns=["step_*", "epoch_*"],
-            )
-
     accelerator.end_training()
 
 if __name__ == "__main__":

trainer.py

@@ -8,9 +8,9 @@ import shutil
 import subprocess
 
 import gradio as gr
-import PIL.Image
 import slugify
 import torch
+from PIL import Image
 from huggingface_hub import HfApi
 
 from app_upload import ModelUploader
@@ -24,9 +24,9 @@ class Trainer:
         self.api = HfApi(token=hf_token)
         self.model_uploader = ModelUploader(hf_token)
 
-    def prepare_dataset(self, reference_images: list, resolution: int,
-                        target_image: PIL.Image, target_mask: PIL.Image,
-                        train_data_dir: pathlib.Path) -> None:
+    def prepare_dataset(self, reference_images: list,
+                        target_image: Image.Image, target_mask: Image.Image,
+                        train_data_dir: pathlib.Path, output_dir: pathlib.Path) -> None:
         shutil.rmtree(train_data_dir, ignore_errors=True)
         train_data_dir.mkdir(parents=True)
         
@@ -34,20 +34,24 @@ class Trainer:
         (train_data_dir / 'target').mkdir(parents=True)
 
         for i, temp_path in enumerate(reference_images):
-            image = PIL.Image.open(temp_path.name)
+            image = Image.open(temp_path.name)
             image = image.convert('RGB')
             out_path = train_data_dir / 'ref' / f'{i:03d}.jpg'
             image.save(out_path, format='JPEG', quality=100)
 
-        target_image = PIL.Image.open(target_image[0].name)
+        target_image = Image.open(target_image[0].name)
         target_image = target_image.convert('RGB')
         out_path = train_data_dir / 'target' / f'target.jpg'
         target_image.save(out_path, format='JPEG', quality=100)
+        out_path = output_dir / f'target.jpg'
+        target_image.save(out_path, format='JPEG', quality=100)
 
-        target_mask = PIL.Image.open(target_mask[0].name)
+        target_mask = Image.open(target_mask[0].name)
         target_mask = target_mask.convert('L')
         out_path = train_data_dir / 'target' / f'mask.jpg'
         target_mask.save(out_path, format='JPEG', quality=100)
+        out_path = output_dir / f'mask.jpg'
+        target_image.save(out_path, format='JPEG', quality=100)
 
     def join_library_org(self) -> None:
         subprocess.run(
@@ -58,8 +62,8 @@ class Trainer:
     def run(
         self,
         reference_images: list | None,
-        target_image: PIL.ImageFile | None,
-        target_mask: PIL.ImageFile | None,
+        target_image: Image.Image | None,
+        target_mask: Image.Image | None,
         output_model_name: str,
         overwrite_existing_model: bool,
         base_model: str,
@@ -104,7 +108,7 @@ class Trainer:
         output_dir.mkdir(parents=True)
 
         train_data_dir = repo_dir / 'training_data' / output_model_name
-        self.prepare_dataset(reference_images, resolution, target_image, target_mask, train_data_dir)
+        self.prepare_dataset(reference_images, target_image, target_mask, train_data_dir, output_dir)
 
         if upload_to_hub:
             self.join_library_org()
@@ -143,8 +147,8 @@ class Trainer:
         subprocess.run(shlex.split(command))
         save_model_card(save_dir=output_dir,
                         base_model=base_model,
-                        target_image=train_data_dir / 'target' / 'target.jpg',
-                        target_mask=train_data_dir / 'target' / 'mask.jpg')
+                        target_image=output_dir / 'target.jpg',
+                        target_mask=output_dir / 'mask.jpg')
 
         message = 'Training completed!'
         print(message)