Refactor and improved logging

app.py

@@ -34,23 +34,25 @@ def generate(
         progress((step_index + 1, pipe.num_timesteps))
         return callback_kwargs
 
-    if model == "timbrooks/instruct-pix2pix":
-        pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model).to("cuda")
+    pipeline_type = (
+        StableDiffusionInstructPix2PixPipeline
+        if model == "timbrooks/instruct-pix2pix"
+        else AutoPipelineForImage2Image
+    )
 
-        images = pipe(
-            prompt=prompt,
-            image=init_image,
-            callback_on_step_end=progress_callback,
-        ).images
-    else:
-        pipe = AutoPipelineForImage2Image.from_pretrained(model).to("cuda")
+    logger.debug(f"Loading pipeline: {dict(model=model)}")
+    pipe = pipeline_type.from_pretrained(model).to("cuda")
 
-        images = pipe(
-            prompt=prompt,
-            image=init_image,
-            strength=strength,
-            callback_on_step_end=progress_callback,
-        ).images
+    logger.debug(f"Generating image: {dict(prompt=prompt)}")
+    additional_args = (
+        {} if model == "timbrooks/instruct-pix2pix" else dict(strength=strength)
+    )
+    images = pipe(
+        prompt=prompt,
+        image=init_image,
+        callback_on_step_end=progress_callback,
+        **additional_args,
+    ).images
     return images[0]