Update app.py

app.py

@@ -6,7 +6,6 @@ import cv2
 import os
 
 from typing import Dict
-from torchvision.transforms.functional import to_tensor, center_crop, resize
 from PIL import Image
 from huggingface_hub import Repository
 
@@ -14,7 +13,8 @@ token = os.environ['model_fetch']
 
 engine_repo = Repository(local_dir="engine", clone_from="felixrosberg/EngageDiffusion", use_auth_token=token)
 
-from engine.ui_model import fetch_model, process_sketch, process_mask, run_model
+from engine.ui_model import fetch_model, run_model
+from engine.ui_gradio import fetch_ui
 
 engage_logo = Image.open("engage_studios_logo.png").resize((700, 88), Image.Resampling.BICUBIC)
 engage_logo_mask = np.array(engage_logo.split()[-1])[..., None] / 255
@@ -103,133 +103,5 @@ def inference(user_state, condition_image, settings, prompt, neg_prompt, inferen
     return user_state["IMAGE_GALLERY"], user_state
 
 
-theme = gr.themes.Base(
-    primary_hue="neutral",
-    radius_size="none",
-).set(
-    body_text_color_dark='*neutral_800',
-    embed_radius='*radius_xxs',
-    button_primary_background_fill='*primary_700',
-    button_primary_background_fill_hover='*primary_400',
-    button_primary_background_fill_hover_dark='*primary_400',
-    button_primary_border_color_dark='*primary_200',
-    button_primary_text_color='*primary_50',
-    button_primary_text_color_dark='*primary_50',
-    button_primary_text_color_hover='*primary_50'
-)
-with gr.Blocks(theme=theme) as engage_automotive_lora_demo:
-    session_state = gr.State(value={"IMAGE_GALLERY": [],
-                                    "SELECTED_IMAGE": None
-                                    })
-    diffused_image_out = gr.Gallery(label='Results', show_label=False,
-                                    columns=[3], rows=[1], object_fit="contain", height="auto",
-                                    format="png")
-    with gr.Group():
-        with gr.Row():
-            prompt_box = gr.Textbox("futuristic dark red car in a white studio",
-                                    label='Prompt')
-            generate_button = gr.Button("Generate", scale=0)
-        with gr.Row():
-            settings_dropdown = gr.Dropdown(
-                ["Text to image", "From sketch", "Inpaint", "Inpaint sketch"], value="Text to image",
-                label="Mode", info="Text to image, prompt only. "
-                                   "From sketch, upload an initial image / sketch in the image editor. "
-                                   "Inpaint sketch, edits the chosen area of an image. Uses the initial "
-                                   "image as base for sketches."
-            )
-        with gr.Accordion("Image Editor", open=False):
-            condition_image = gr.ImageEditor(type='pil', show_label=False,
-                                             brush=gr.Brush(colors=["#000000"], color_mode="fixed"))
-        with gr.Row():
-            with gr.Accordion("Settings", open=False):
-                neg_prompt_box = gr.Textbox(
-                    "blurry, poor quality, unrealistic",
-                    label='Negative Prompt')
-                seed_box = gr.Number(42, label='Seed')
-                inference_steps = gr.Slider(0, 20, value=8,
-                                            label='Inference Steps', step=1)
-                num_images = gr.Slider(1, 3, value=2, label='Number of Images', step=1)
-                guidance_scale = gr.Slider(0, 10, value=1.5,
-                                           label='Guidance Scale', step=0.1)
-                guidance_rescale = gr.Slider(0.0, 1.0, value=0.0,
-                                             label='Guidance Rescale', step=0.1)
-                height = gr.Slider(128, 2048, value=1024, label='Image Height', step=64)
-                width = gr.Slider(128, 2048, value=1024, label='Image Width', step=64)
-                condition_influence = gr.Slider(0.0, 1.0, value=0.5, label='Condition Influence')
-                sketch_detail = gr.Slider(0.0, 1.0, value=0.5, label='Sketch Detail')
-                sketch_softness = gr.Slider(0.0, 1.0, value=0.5, label='Sketch Softness')
-                inpaint_strength = gr.Slider(0.0, 1.0, value=0.8, label='Inpaint Strength')
-                enable_freeu = gr.Checkbox(True, label='FreeU',
-                                           info='Enables FreeU scaling factors.')
-            with gr.Accordion("Stylation (Experimental)", open=False):
-                with gr.Row():
-                    exposure = gr.Slider(-1.0, 1.0, value=0.0, label='Exposure')
-                    enable_stylation = gr.Checkbox(label='Enable Stylation',
-                                                   info='EXPERIMENTAL: We apologize for the ambiguity, '
-                                                        'please play around with the sliders to '
-                                                        'find a style you like!'
-                                                        'Warning: Will slow down the generation time.')
-                with gr.Accordion("Style A - Engage Studios Futuristic", open=False):
-                    style_A_down = gr.Slider(-1.0, 1.0, value=0.0, label='down')
-                    style_A_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid')
-                    style_A_up = gr.Slider(-1.0, 1.0, value=0.0, label='up')
-                with gr.Accordion("Style B - Lighting", open=False):
-                    style_B_down = gr.Slider(-1.0, 1.0, value=0.0, label='down')
-                    style_B_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid')
-                    style_B_up = gr.Slider(-1.0, 1.0, value=0.0, label='up')
-                with gr.Accordion("Style C - Details A", open=False):
-                    style_C_down = gr.Slider(-1.0, 1.0, value=0.0, label='down')
-                    style_C_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid')
-                    style_C_up = gr.Slider(-1.0, 1.0, value=0.0, label='up')
-                with gr.Accordion("Style D - Details B", open=False):
-                    style_D_down = gr.Slider(-1.0, 1.0, value=0.0, label='down')
-                    style_D_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid')
-                    style_D_up = gr.Slider(-1.0, 1.0, value=0.0, label='up')
-
-    generate_button.click(inference,
-                          inputs=[session_state,
-                                  condition_image,
-                                  settings_dropdown,
-                                  prompt_box,
-                                  neg_prompt_box,
-
-                                  inference_steps,
-                                  num_images,
-
-                                  guidance_scale,
-                                  guidance_rescale,
-                                  enable_freeu,
-
-                                  height,
-                                  width,
-
-                                  condition_influence,
-
-                                  sketch_detail,
-                                  sketch_softness,
-                                  inpaint_strength,
-
-                                  exposure,
-                                  enable_stylation,
-
-                                  style_A_down,
-                                  style_A_mid,
-                                  style_A_up,
-
-                                  style_B_down,
-                                  style_B_mid,
-                                  style_B_up,
-
-                                  style_C_down,
-                                  style_C_mid,
-                                  style_C_up,
-
-                                  style_D_down,
-                                  style_D_mid,
-                                  style_D_up,
-
-                                  seed_box],
-                          outputs=[diffused_image_out, session_state],
-                          show_progress=True)
-
-engage_automotive_lora_demo.launch()
+engage_demo = fetch_ui(inference)
+engage_demo.launch()