Update app.py

app.py

@@ -110,21 +110,43 @@ controlnet_models = {
     "scribble": "lllyasviel/sd-controlnet-scribble"
 }
 
-def infer(**kwargs):
-    generator = torch.Generator(device).manual_seed(kwargs['seed'])
+def infer(
+    prompt, 
+    negative_prompt, 
+    width=512, 
+    height=512, 
+    num_inference_steps=20, 
+    model_id='CompVis/stable-diffusion-v1-4', 
+    seed=42, 
+    guidance_scale=7.0, 
+    lora_scale=0.5,
+    cn_enable=False,
+    cn_strength=0.0,
+    cn_mode='edge_detection',
+    cn_image=None,
+    ip_enable=False,
+    ip_scale=0.5,
+    ip_image=None,
+    progress=gr.Progress(track_tqdm=True)
+    ):
+
+    generator = torch.Generator(device).manual_seed(seed)
     
     with torch_inference_mode():
         pipe = pipeline_mgr.initialize_pipeline(
-            kwargs['model_id'], 
-            controlnet_models.get(kwargs['cn_mode'], controlnet_models['edge_detection'])
+            model_id, 
+            controlnet_models.get(cn_mode, controlnet_models['edge_detection'])
         )
         
-        if kwargs['cn_enable'] and not kwargs['cn_image']:
+        if cn_enable and not cn_image:
             raise gr.Error("ControlNet enabled but no image provided!")
+
+        if ip_enable and not ip_image:
+            raise gr.Error("IP-Adapter enabled but no image provided!")
         
         prompt_emb, negative_emb = process_embeddings(
-            kwargs['prompt'], 
-            kwargs['negative_prompt'], 
+            prompt, 
+            negative_prompt, 
             pipe.tokenizer, 
             pipe.text_encoder
         )
@@ -132,27 +154,27 @@ def infer(**kwargs):
         params = {
             'prompt_embeds': prompt_emb,
             'negative_prompt_embeds': negative_emb,
-            'guidance_scale': kwargs['guidance_scale'],
-            'num_inference_steps': kwargs['num_inference_steps'],
-            'width': kwargs['width'],
-            'height': kwargs['height'],
+            'guidance_scale': guidance_scale,
+            'num_inference_steps': num_inference_steps,
+            'width': width,
+            'height': height,
             'generator': generator
         }
         
-        if kwargs['cn_enable']:
+        if cn_enable:
             params['image'] = process_control_image(
-                kwargs['cn_image'],
-                kwargs['cn_mode'],
+                cn_image,
+                cn_mode,
                 pipeline_mgr.get_hed_detector()
             )
-            params['controlnet_conditioning_scale'] = kwargs['cn_strength']
+            params['controlnet_conditioning_scale'] = cn_strength
             
-        if kwargs.get('ip_enable', False):
+        if ip_enable:
             pipe.load_ip_adapter(IP_ADAPTER, subfolder="models", weight_name=IP_ADAPTER_WEIGHT_NAME)
-            params['ip_adapter_image'] = load_image(kwargs['ip_image']).convert('RGB')
-            pipe.set_ip_adapter_scale(kwargs.get('ip_scale', 0.5))
+            params['ip_adapter_image'] = load_image(ip_image).convert('RGB')
+            pipe.set_ip_adapter_scale(ip_scale)
             
-        pipe.fuse_lora(lora_scale=kwargs.get('lora_scale', 0.5))
+        pipe.fuse_lora(lora_scale=lora_scale)
         
         return pipe(**params).images[0]