Add app.py

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+import spaces
+from diffusers import AutoPipelineForImage2Image
+from loguru import logger
+from PIL import Image
+
+models = [
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    "stabilityai/sdxl-turbo",
+    "timbrooks/instruct-pix2pix",
+]
+
+
+@spaces.GPU
+@logger.catch(reraise=True)
+def generate(
+    model: str,
+    prompt: str,
+    init_image: Image.Image,
+    strength: float,
+    progress=gr.Progress(),
+):
+    logger.info(
+        f"Starting image generation: {dict(model=model, prompt=prompt, image=init_image, strength=strength)}"
+    )
+
+    pipe = AutoPipelineForImage2Image.from_pretrained(model).to("cuda")
+
+    # Downscale the image
+    init_image.thumbnail((1024, 1024))
+
+    def progress_callback(pipe, step_index, timestep, callback_kwargs):
+        logger.debug(
+            f"Callback: {dict(num_timesteps=pipe.num_timesteps, step_index=step_index, timestep=timestep)}"
+        )
+        progress((step_index + 1, pipe.num_timesteps))
+        return callback_kwargs
+
+    images = pipe(
+        prompt=prompt,
+        image=init_image,
+        strength=strength,
+        callback_on_step_end=progress_callback,
+    ).images
+    return images[0]
+
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Dropdown(
+            label="Model", choices=models, value=models[0], allow_custom_value=True
+        ),
+        gr.Text(label="Prompt"),
+        gr.Image(label="Init image", type="pil"),
+        gr.Slider(label="Strength", minimum=0, maximum=1, value=0.3),
+    ],
+    outputs=[gr.Image(label="Output")],
+)
+
+demo.launch()