Update app.py

app.py

@@ -14,7 +14,7 @@ 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
+from engine.ui_model import fetch_model, process_sketch, process_mask, run_model, run_text_to_image
 
 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
@@ -24,35 +24,7 @@ pipe = fetch_model()
 pipe.to('cuda')
 
 
-@spaces.GPU
-def run_text_to_image(prompt=None, neg_prompt=None,
-                      inference_steps=8, num_images=2,
-                      guidance_scale=2.0,
-                      guidance_rescale=0.0,
-
-                      height=1024,
-                      width=1024,
-
-                      condition_scale=0.5,
-
-                      exposure=0.0,
-                      progress=gr.Progress()):
-    images = pipe(prompt=prompt,
-                  negative_prompt=neg_prompt,
-                  num_images_per_prompt=num_images,
-                  num_inference_steps=inference_steps,
-                  height=height,
-                  width=width,
-                  guidance_scale=guidance_scale,
-                  guidance_rescale=guidance_rescale,
-                  controlnet_conditioning_scale=condition_scale,
-                  gradio_progress=progress,
-                  cross_attention_kwargs={"scale": exposure}
-                  ).images
-    return images
-
-
-def run_model(user_state, condition_image, settings, prompt, neg_prompt, inference_steps=8, num_images=2,
+def inference(user_state, condition_image, settings, prompt, neg_prompt, inference_steps=8, num_images=2,
               guidance_scale=2.0,
               guidance_rescale=0.0,
               enable_freeu=False,
@@ -86,47 +58,41 @@ def run_model(user_state, condition_image, settings, prompt, neg_prompt, inferen
 
               seed=None,
               progress=gr.Progress()):
-    # prompt += ", shot with a mirrorless, 35mm, photography, real, 8k, photorealistic, "
-    prompt += "best quality, HD, ~*~aesthetic~*~"
-
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-    progress(0, desc="Thinking...", total=int(inference_steps))
-
-    if enable_freeu:
-        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
-    else:
-        pipe.disable_freeu()
-
-    if exposure != 0.0 and enable_stylation:
-        pipe.enable_lora()
-        adapter_weight_scales_ENGAGE = {"unet": {"down": style_1_down, "mid": style_1_mid, "up": style_1_up}}
-        adapter_weight_scales_FILM = {"unet": {"down": style_2_down, "mid": style_2_mid, "up": style_2_up}}
-        adapter_weight_scales_MJ = {"unet": {"down": style_3_down, "mid": style_3_mid, "up": style_3_up}}
-        adapter_weight_scales_MORE_ART = {"unet": {"down": style_4_down, "mid": style_4_mid, "up": style_4_up}}
-
-        pipe.set_adapters(["ENGAGE_LORA", "FILM_LORA", "MJ_LORA", "MORE_ART_LORA"],
-                          [adapter_weight_scales_ENGAGE,
-                           adapter_weight_scales_FILM,
-                           adapter_weight_scales_MJ,
-                           adapter_weight_scales_MORE_ART])
-    else:
-        pipe.set_adapters(["ENGAGE_LORA", "FILM_LORA", "MJ_LORA", "MORE_ART_LORA"],
-                          adapter_weights=[0.0, 0.0, 0.0, 0.0])
-        pipe.disable_lora()
-
-    images = run_text_to_image(prompt=prompt,
-                               neg_prompt=neg_prompt,
-                               num_images=num_images,
-                               inference_steps=inference_steps,
-                               height=height,
-                               width=width,
-                               guidance_scale=guidance_scale,
-                               guidance_rescale=guidance_rescale,
-                               condition_scale=condition_scale,
-                               progress=progress,
-                               exposure=exposure)
+
+    images = run_model(pipe, user_state, condition_image, settings, prompt, neg_prompt, inference_steps=8, num_images=2,
+              guidance_scale,
+              guidance_rescale,
+              enable_freeu,
+
+              height,
+              width,
+
+              condition_scale,
+              sketch_detail,
+              sketch_softness,
+              inpaint_strength,
+
+              exposure,
+              enable_stylation,
+
+              style_1_down,
+              style_1_mid,
+              style_1_up,
+
+              style_2_down,
+              style_2_mid,
+              style_2_up,
+
+              style_3_down,
+              style_3_mid,
+              style_3_up,
+
+              style_4_down,
+              style_4_mid,
+              style_4_up,
+
+              seed,
+              progress)
 
     for idx, im in enumerate(images):
         im = np.asarray(im).copy()