update

app.py

@@ -24,9 +24,6 @@ import gradio as gr
 import torch
 
 from inference import inference_fn
-# from inference_custom_diffusion import InferencePipeline
-# from trainer import Trainer
-# from uploader import upload
 
 
 # def parse_args() -> argparse.Namespace:
@@ -49,9 +46,12 @@ It is recommended to upgrade to GPU in Settings after duplicating this space to
 DETAILDESCRIPTION='''
 ReVersion
 '''
-# DETAILDESCRIPTION='''
-# ReVersion: <R> represents the learned text token for a relation. Use <R> in your prompt for relation-specific generation.
-# '''
+DETAILDESCRIPTION='''
+ReVersion: <R> represents the learned text token for a relation. Use <R> in your prompt for relation-specific generation.
+<center>
+<img src="teaser.jpg" width="600" align="center">
+</center>
+'''
 # DETAILDESCRIPTION='''
 # Custom Diffusion allows you to fine-tune text-to-image diffusion models, such as Stable Diffusion, given a few images of a new concept (~4-20).
 # We fine-tune only a subset of model parameters, namely key and value projection matrices, in the cross-attention layers and the modifier token used to represent the object.
@@ -124,16 +124,16 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
                 #     placeholder='Example: "<R>"')
 
                 with gr.Accordion('Other Parameters', open=False):
+                    num_samples = gr.Slider(label='Number of Images to Generate',
+                                               minimum=4,
+                                               maximum=8,
+                                               step=2,
+                                               value=6)
                     guidance_scale = gr.Slider(label='Classifier-Free Guidance Scale',
                                                minimum=0,
                                                maximum=50,
                                                step=0.1,
                                                value=7.5)
-                    num_samples = gr.Slider(label='Number of Images to Generate',
-                                               minimum=0,
-                                               maximum=10.,
-                                               step=1,
-                                               value=10)
                     ddim_steps = gr.Slider(label='Number of DDIM Sampling Steps',
                                                minimum=10,
                                                maximum=100,

inference.py

@@ -54,12 +54,6 @@ def inference_fn(
         pipe = StableDiffusionPipeline.from_pretrained(os.path.join('experiments', model_id),torch_dtype=torch.float16).to('cuda')
     else:
         pipe = StableDiffusionPipeline.from_pretrained(os.path.join('experiments', model_id)).to('cpu')
-    # # make directory to save images
-    # image_root_folder = os.path.join('experiments', model_id, 'inference')
-    # os.makedirs(image_root_folder, exist_ok = True)
-
-    # if prompt is None and args.template_name is None:
-    #     raise ValueError("please input a single prompt through'--prompt' or select a batch of prompts using '--template_name'.")
 
     # single text prompt
     if prompt is not None:
@@ -67,33 +61,17 @@ def inference_fn(
     else:
         prompt_list = []
 
-    # if args.template_name is not None:
-    #     # read the selected text prompts for generation
-    #     prompt_list.extend(inference_templates[args.template_name])
-
     for prompt in prompt_list:
         # insert relation prompt <R>
         # prompt = prompt.lower().replace("<r>", "<R>").format(placeholder_string)
         prompt = prompt.lower().replace("<r>", "<R>").format("<R>")
 
-
-        # # make sub-folder
-        # image_folder = os.path.join(image_root_folder, prompt, 'samples')
-        # os.makedirs(image_folder, exist_ok = True)
-
         # batch generation
         images = pipe(prompt, num_inference_steps=ddim_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_samples).images
 
-        # # save generated images
-        # for idx, image in enumerate(images):
-        #     image_name = f"{str(idx).zfill(4)}.png"
-        #     image_path = os.path.join(image_folder, image_name)
-        #     image.save(image_path)
-
         # save a grid of images
         image_grid = make_image_grid(images, rows=2, cols=math.ceil(num_samples/2))
         print(image_grid)
-        # image_grid_path = os.path.join(image_root_folder, prompt, f'{prompt}.png')
 
         return image_grid