Fix pickling issue with ZeroGPU tuning

app.py

@@ -13,14 +13,51 @@ models = [
 DEFAULT_MODEL = "stabilityai/stable-diffusion-xl-refiner-1.0"
 
 
+def generate_image(
+    model: str,
+    prompt: str,
+    init_image: Image.Image,
+    strength: float,
+    progress,
+):
+    pipeline_type = (
+        StableDiffusionInstructPix2PixPipeline
+        if model == "timbrooks/instruct-pix2pix"
+        else AutoPipelineForImage2Image
+    )
+
+    logger.debug(f"Loading pipeline: {dict(model=model)}")
+    pipe = pipeline_type.from_pretrained(model).to("cuda")
+
+    logger.debug(f"Generating image: {dict(prompt=prompt)}")
+    additional_args = (
+        {} if model == "timbrooks/instruct-pix2pix" else dict(strength=strength)
+    )
+
+    def progress_callback(pipe, step_index, timestep, callback_kwargs):
+        logger.trace(
+            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,
+        callback_on_step_end=progress_callback,
+        **additional_args,
+    ).images
+    return images[0]
+
+
 @spaces.GPU
-def gpu(fn):
-    return fn()
+def gpu(*args, **kwargs):
+    return generate_image(*args, **kwargs)
 
 
 @spaces.GPU(duration=180)
-def gpu_3min(fn):
-    return fn()
+def gpu_3min(*args, **kwargs):
+    return generate_image(*args, **kwargs)
 
 
 @logger.catch(reraise=True)
@@ -38,38 +75,9 @@ def generate(
     # Downscale the image
     init_image.thumbnail((1024, 1024))
 
-    def progress_callback(pipe, step_index, timestep, callback_kwargs):
-        logger.trace(
-            f"Callback: {dict(num_timesteps=pipe.num_timesteps, step_index=step_index, timestep=timestep)}"
-        )
-        progress((step_index + 1, pipe.num_timesteps))
-        return callback_kwargs
-
-    pipeline_type = (
-        StableDiffusionInstructPix2PixPipeline
-        if model == "timbrooks/instruct-pix2pix"
-        else AutoPipelineForImage2Image
-    )
-
-    logger.debug(f"Loading pipeline: {dict(model=model)}")
-    pipe = pipeline_type.from_pretrained(model)
-
-    logger.debug(f"Generating image: {dict(prompt=prompt)}")
-    additional_args = (
-        {} if model == "timbrooks/instruct-pix2pix" else dict(strength=strength)
-    )
-
     gpu_runner = gpu_3min if model == "timbrooks/instruct-pix2pix" else gpu
 
-    images = gpu_runner(
-        lambda: pipe.to("cuda")(
-            prompt=prompt,
-            image=init_image,
-            callback_on_step_end=progress_callback,
-            **additional_args,
-        ).images
-    )
-    return images[0]
+    return gpu_runner(model, prompt, init_image, strength, progress)
 
 
 demo = gr.Interface(