add style guidance

app.py

@@ -19,9 +19,16 @@ if disable_safety:
         return images, False
     pipe.safety_checker = null_safety
 
+
+style_embeddings = {
+    'Concept Art': torch.load('embeddings/concept_art.pt'), 'Animation': torch.load('embeddings/animation.pt'), 'Character Design': torch.load('embeddings/character_design.pt')
+    , 'Portrait Photo': torch.load('embeddings/portrait_photo.pt'), 'Architecture': torch.load('embeddings/architecture.pt')
+}
     
 def infer(prompt, steps, scale, seed, editing_prompt_1 = None, reverse_editing_direction_1 = False, edit_warmup_steps_1=10, edit_guidance_scale_1=5, edit_threshold_1=0.95, 
           editing_prompt_2 = None, reverse_editing_direction_2 = False, edit_warmup_steps_2=10, edit_guidance_scale_2=5, edit_threshold_2=0.95, 
+          edit_style=None,
+          reverse_editing_direction_style = False, edit_warmup_steps_style=5, edit_guidance_scale_style=7, edit_threshold_style=0.8, 
           edit_momentum_scale=0.5, edit_mom_beta=0.6):
     
     
@@ -42,15 +49,52 @@ def infer(prompt, steps, scale, seed, editing_prompt_1 = None, reverse_editing_d
         del edit_warmup_steps[index]
         del edit_guidance_scale[index]
         del edit_threshold[index]
+    editing_prompt_embeddings = None
+    
     
+    if edit_style is not None:
+        editing_prompt = None
+        reverse_editing_direction = reverse_editing_direction_style
+        edit_warmup_steps = edit_warmup_steps_style
+        edit_guidance_scale = edit_guidance_scale_style
+        edit_threshold = edit_threshold_style
+        editing_prompt_embeddings = style_embeddings[edit_style]
     
     gen.manual_seed(seed)
     images.extend(pipe(prompt, guidance_scale=scale, num_inference_steps=steps, generator=gen,
-                      editing_prompt=editing_prompt, reverse_editing_direction=reverse_editing_direction, edit_warmup_steps=edit_warmup_steps, edit_guidance_scale=edit_guidance_scale, 
+                      editing_prompt=editing_prompt, editing_prompt_embeddings=editing_prompt_embeddings,
+                       reverse_editing_direction=reverse_editing_direction, edit_warmup_steps=edit_warmup_steps, edit_guidance_scale=edit_guidance_scale,                                            
                         edit_momentum_scale=edit_momentum_scale, edit_mom_beta=edit_mom_beta
                       ).images)
 
-    return images
+    return zip(images, ['Original', edit_style if edit_style is not None else 'SEGA'])
+
+def reset_style():
+    radio = gr.Radio(label='Style', choices=['Concept Art', 'Animation', 'Character Design', 'Portrait Photo', 'Architecture'])
+    return radio
+
+def reset_text():
+    text_1 = gr.Textbox(
+                        label="Edit Prompt 1",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your 1st edit prompt",
+                    ).style(
+                        border=(True, False, True, True),
+                        rounded=(True, False, False, True),
+                        container=False,
+                    )
+    text_2 = gr.Textbox(
+                        label="Edit Prompt 2",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your 2nd edit prompt",
+                    ).style(
+                        border=(True, False, True, True),
+                        rounded=(True, False, False, True),
+                        container=False,
+                    )
+    return text_1, text_2
 
 css = """
         a {
@@ -144,13 +188,18 @@ examples = [
         'sunglasses',
         False,
         10,
-        6,
+        5,
         0.95,
         '',
         False,
         10,
         5,
-        0.95
+        0.95,
+        '',
+        False,
+        5,
+        7,
+        0.8,
     ],
     [
         'an image of a crowded boulevard, realistic, 4k',
@@ -166,7 +215,12 @@ examples = [
         False,
         10,
         5,
-        0.95
+        0.95,
+        '',
+        False,
+        5,
+        7,
+        0.8
     ],
     [
         'a castle next to a river',
@@ -182,6 +236,11 @@ examples = [
         False,
         18,
         6,
+        0.8,
+        '',
+        False,
+        5,
+        7,
         0.8
     ],
     [
@@ -198,7 +257,12 @@ examples = [
         False,
         5,
         5,
-        0.9
+        0.9,
+        '',
+        False,
+        5,
+        7,
+        0.8
     ],
     [
         'a photo of a flowerpot',
@@ -214,7 +278,12 @@ examples = [
         False,
         10,
         5,
-        0.95
+        0.95,
+        '',
+        False,
+        5,
+        7,
+        0.8
     ],
      [
         'a photo of the face of a woman',
@@ -230,10 +299,79 @@ examples = [
         False,
         13,
         3,
-        0.925
+        0.925,
+        '',
+        False,
+        5,
+        7,
+        0.8
+    ],
+     [
+        'temple in ruines, forest, stairs, columns',
+        50,
+        7,
+        11,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        'Animation',
+        False,
+        5,
+        7,
+        0.8
+    ],
+    [
+        'city made out of glass',
+        50,
+        7,
+        16,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        'Concept Art',
+        False,
+        10,
+        8,
+        0.8
+    ],
+     [
+        'a man riding a horse',
+        50,
+        7,
+        11,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        '',
+        False,
+        10,
+        5,
+        0.95,
+        'Character Design',
+        False,
+        11,
+        8,
+        0.9
     ],
 ]
 
+
 with block:
     gr.HTML(
         """
@@ -268,48 +406,58 @@ with block:
                     margin=False,
                     rounded=(False, True, True, False),
                 )
-            with gr.Box():
-                with gr.Row().style(mobile_collapse=False, equal_height=True):
-                    edit_1 = gr.Textbox(
-                        label="Edit Prompt 1",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your 1st edit prompt",
-                    ).style(
-                        border=(True, False, True, True),
-                        rounded=(True, False, False, True),
-                        container=False,
-                    )
-                    with gr.Group():
-                        with gr.Row().style(mobile_collapse=False, equal_height=True):
-                            rev_1 = gr.Checkbox(
-                                label='Reverse')
-                            warmup_1 = gr.Slider(label='Warmup', minimum=0, maximum=50, value=10, step=1, interactive=True)
-                            scale_1 = gr.Slider(label='Scale', minimum=1, maximum=10, value=5, step=0.25, interactive=True)
-                            threshold_1 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, value=0.95, steps=0.01, interactive=True)
-                with gr.Row().style(mobile_collapse=False, equal_height=True):
-                    edit_2 = gr.Textbox(
-                        label="Edit Prompt 2",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your 2nd edit prompt",
-                    ).style(
-                        border=(True, False, True, True),
-                        rounded=(True, False, False, True),
-                        container=False,
-                    )
-                    with gr.Group():
-                        with gr.Row().style(mobile_collapse=False, equal_height=True):
-                            rev_2 = gr.Checkbox(
-                                label='Reverse')
-                            warmup_2 = gr.Slider(label='Warmup', minimum=0, maximum=50, value=10, step=1, interactive=True)
-                            scale_2 = gr.Slider(label='Scale', minimum=1, maximum=10, value=5, step=0.25, interactive=True)
-                            threshold_2 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, value=0.95, steps=0.01, interactive=True)
-            
-            
-
+            with gr.Tabs() as tabs:
+                with gr.TabItem('Text Guidance', id=0):
+                    with gr.Row().style(mobile_collapse=False, equal_height=True):
+                        edit_1 = gr.Textbox(
+                            label="Edit Prompt 1",
+                            show_label=False,
+                            max_lines=1,
+                            placeholder="Enter your 1st edit prompt",
+                        ).style(
+                            border=(True, False, True, True),
+                            rounded=(True, False, False, True),
+                            container=False,
+                        )
+                        with gr.Group():
+                            with gr.Row().style(mobile_collapse=False, equal_height=True):
+                                rev_1 = gr.Checkbox(
+                                    label='Negative Guidance')
+                                warmup_1 = gr.Slider(label='Warmup', minimum=0, maximum=50, value=10, step=1, interactive=True)
+                                scale_1 = gr.Slider(label='Scale', minimum=1, maximum=10, value=5, step=0.25, interactive=True)
+                                threshold_1 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, value=0.95, steps=0.01, interactive=True)
+                    with gr.Row().style(mobile_collapse=False, equal_height=True):
+                        edit_2 = gr.Textbox(
+                            label="Edit Prompt 2",
+                            show_label=False,
+                            max_lines=1,
+                            placeholder="Enter your 2nd edit prompt",
+                        ).style(
+                            border=(True, False, True, True),
+                            rounded=(True, False, False, True),
+                            container=False,
+                        )
+                        with gr.Group():
+                            with gr.Row().style(mobile_collapse=False, equal_height=True):
+                                rev_2 = gr.Checkbox(
+                                    label='Negative Guidance')
+                                warmup_2 = gr.Slider(label='Warmup', minimum=0, maximum=50, value=10, step=1, interactive=True)
+                                scale_2 = gr.Slider(label='Scale', minimum=1, maximum=10, value=5, step=0.25, interactive=True)
+                                threshold_2 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, value=0.95, steps=0.01, interactive=True)
+                with gr.TabItem("Style Guidance", id=1):
+                    with gr.Row().style(mobile_collapse=False, equal_height=True):
+                        style = gr.Radio(label='Style', choices=['Concept Art', 'Animation', 'Character Design', 'Portrait Photo', 'Architecture'], interactive=True)
+                        with gr.Group():
+                            with gr.Row().style(mobile_collapse=False, equal_height=True):
+                                rev_style = gr.Checkbox(
+                                    label='Negative Guidance', interactive=False)
+                                warmup_style = gr.Slider(label='Warmup', minimum=0, maximum=50, value=5, step=1, interactive=True)
+                                scale_style = gr.Slider(label='Scale', minimum=1, maximum=10, value=7, step=0.25, interactive=True)
+                                threshold_style = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, value=0.8, steps=0.01, interactive=True)
+        
+        
         gallery = gr.Gallery(
-            label="Generated images", show_label=False, elem_id="gallery"
+            label=("Generated images"), show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
 
@@ -324,13 +472,29 @@ with block:
                 #randomize=True,
             )
 
+        
+        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2, style, rev_style, warmup_style, scale_style, threshold_style], outputs=gallery, cache_examples=False)
+        ex.dataset.headers = ['Prompt', 'Steps', 'Scale', 'Seed', 'Edit Prompt 1', 'Negation 1', 'Warmup 1', 'Scale 1', 'Threshold 1', 'Edit Prompt 2', 'Negation 2', 'Warmup 2', 'Scale 2', 'Threshold 2', 'Style', 'Style Negation', 'Style Warmup', 'Style Scale', 'Style Threshold']
 
-        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2], outputs=gallery, cache_examples=False)
-        ex.dataset.headers = ['Prompt', 'Steps', 'Scale', 'Seed', 'Edit Prompt 1', 'Reverse 1', 'Warmup 1', 'Scale 1', 'Threshold 1', 'Edit Prompt 2', 'Reverse 2', 'Warmup 2', 'Scale 2', 'Threshold 2']
 
+        text.submit(infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2, style, rev_style, warmup_style, scale_style, threshold_style], outputs=gallery)
+        btn.click(infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2, style, rev_style, warmup_style, scale_style, threshold_style], outputs=gallery)
+        #btn.click(change_tab, None, tabs)
+        
+        edit_1.change(reset_style, outputs=style)
+        edit_2.change(reset_style, outputs=style)
+        
+        rev_1.change(reset_style, outputs=style)
+        rev_2.change(reset_style, outputs=style)
+        
+        warmup_1.change(reset_style, outputs=style)
+        warmup_2.change(reset_style, outputs=style)
+        
+        threshold_1.change(reset_style, outputs=style)
+        threshold_2.change(reset_style, outputs=style)
+        #style.change(reset_text, outputs=[edit_1, edit_2])
 
-        text.submit(infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2], outputs=gallery)
-        btn.click(infer, inputs=[text, steps, scale, seed, edit_1, rev_1, warmup_1, scale_1, threshold_1, edit_2, rev_2, warmup_2, scale_2, threshold_2], outputs=gallery)
+        
         gr.HTML(
             """
                 <div class="footer">