img2img

app.py

@@ -15,18 +15,34 @@ MAX_IMAGE_SIZE = 2048
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
 @spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-        prompt=prompt, 
-        width=width,
-        height=height,
-        num_inference_steps=num_inference_steps, 
-        generator=generator,
-        guidance_scale=0.0
-    ).images[0]
+    
+    if init_image is not None:
+        # Process img2img
+        init_image = pipe.preprocess(init_image).unsqueeze(0).to(device, dtype)
+        image = pipe(
+            prompt=prompt,
+            init_image=init_image,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0
+        ).images[0]
+    else:
+        # Process text2img
+        image = pipe(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0
+        ).images[0]
+    
     return image, seed
 
 # Define example prompts
@@ -91,7 +107,9 @@ with gr.Blocks(css=css) as demo:
             )
             run_button = gr.Button("Run", scale=0)
 
-        result = gr.Image(label="Result", show_label=False)
+        with gr.Row():
+            init_image = gr.Image(label="Initial Image (optional)", type="pil", optional=True)
+            result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
             seed = gr.Slider(
@@ -139,7 +157,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
+        inputs=[prompt, init_image, seed, randomize_seed, width, height, num_inference_steps],
         outputs=[result, seed]
     )