Update app.py

app.py

@@ -45,7 +45,7 @@ style_list = [
     },
 ]
 
-STYLE_NAMES = [style["name"] for style in style_list]
+STYLE_NAMES = [s["name"] for s in style_list]
 DEFAULT_STYLE_NAME = STYLE_NAMES[0]
 
 grid_sizes = {
@@ -79,11 +79,10 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
+    grid_x, grid_y = grid_sizes.get(grid_size, (1, 1))
+    num_images = grid_x * grid_y
 
-    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (1, 1))
-    num_images = grid_size_x * grid_size_y
-
-    options = {
+    opts = {
         "prompt": styled_prompt,
         "negative_prompt": styled_negative_prompt,
         "guidance_scale": guidance_scale,
@@ -94,13 +93,14 @@ def infer(
         "num_images_per_prompt": num_images,
     }
 
-    torch.cuda.empty_cache()  # Clear GPU memory
-    result = pipe(**options)
-
-    grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
+    torch.cuda.empty_cache()
+    res = pipe(**opts)
 
-    for i, img in enumerate(result.images[:num_images]):
-        grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
+    grid_img = Image.new('RGB', (width * grid_x, height * grid_y))
+    for i, img in enumerate(res.images[:num_images]):
+        x = (i % grid_x) * width
+        y = (i // grid_x) * height
+        grid_img.paste(img, (x, y))
 
     return grid_img, seed
 
@@ -111,36 +111,37 @@ examples = [
     "A cat holding a sign that says hello world --ar 85:128 --v 6.0 --style raw"
 ]
 
-#css = '''
-#.gradio-container{max-width: 585px !important}
-#h1{text-align:center}
-#footer {
-#    visibility: hidden
-#}
-#'''
+css = '''
+.gradio-container {
+    max-width: 585px !important;
+    margin: 0 auto !important;
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+    justify-content: center;
+}
+h1 { text-align: center; }
+footer { visibility: hidden; }
+'''
 
-#with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-with gr.Blocks() as demo:
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("## GRID 6X🪨")
+        gr.Markdown("## GRID 6X🪨")
 
         with gr.Row():
             prompt = gr.Text(
-                label="Prompt",
                 show_label=False,
                 max_lines=1,
                 placeholder="Enter your prompt",
                 container=False,
             )
-
             run_button = gr.Button("Run", scale=0, variant="primary")
 
-        result = gr.Image(label="Result", show_label=False)
-
+        result = gr.Image(show_label=False)
 
-        with gr.Row(visible=True):
+        with gr.Row():
             grid_size_selection = gr.Dropdown(
-                choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
+                choices=list(grid_sizes.keys()),
                 value="1x1",
                 label="Grid Size"
             )
@@ -151,82 +152,39 @@ with gr.Blocks() as demo:
                 max_lines=1,
                 placeholder="Enter a negative prompt",
                 value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-                visible=False,
             )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
+            seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
             with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=512,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=512,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
+                width = gr.Slider(512, MAX_IMAGE_SIZE, step=32, value=1024, label="Width")
+                height = gr.Slider(512, MAX_IMAGE_SIZE, step=32, value=1024, label="Height")
 
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=7.5,
-                    step=0.1,
-                    value=0.0,
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=8,
-                )       
-                
-                style_selection = gr.Radio(
-                    show_label=True,
-                    container=True,
-                    interactive=True,
-                    choices=STYLE_NAMES,
-                    value=DEFAULT_STYLE_NAME,
-                    label="Quality Style",
-                )
-
-        gr.Examples(examples=examples, 
-                    inputs=[prompt], 
-                    outputs=[result, seed], 
-                    fn=infer, 
-                    cache_examples=False)
+                guidance_scale = gr.Slider(0.0, 7.5, step=0.1, value=7.5, label="Guidance scale")
+                num_inference_steps = gr.Slider(1, 50, step=1, value=10, label="Number of inference steps")
+
+            style_selection = gr.Radio(
+                choices=STYLE_NAMES,
+                value=DEFAULT_STYLE_NAME,
+                label="Quality Style",
+            )
+
+        gr.Examples(
+            examples=examples,
+            inputs=[prompt],
+            outputs=[result, seed],
+            fn=infer,
+            cache_examples=False
+        )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            style_selection,
-            grid_size_selection,
+            prompt, negative_prompt, seed, randomize_seed,
+            width, height, guidance_scale, num_inference_steps,
+            style_selection, grid_size_selection
         ],
         outputs=[result, seed],
     )