Update flux1_img2img.py

flux1_img2img.py

@@ -4,32 +4,42 @@ from PIL import Image
 import sys
 import spaces
 
-# I only test with FLUX.1-schnell
+# Tested with FLUX.1-schnell
 
 @spaces.GPU
 def process_image(image, mask_image, prompt="a person", model_id="black-forest-labs/FLUX.1-schnell", strength=0.75, seed=0, num_inference_steps=4):
-    print("start process image process_image")
+    print("Starting process_image")
     if image is None:
-        print("empty input image returned")
+        print("Empty input image returned.")
         return None
 
-    # Ensure image is in RGB mode (helps with WebP and other formats)
+    # Ensure the image is in RGB mode (this handles formats like WebP and JFIF)
     if image.mode != "RGB":
         image = image.convert("RGB")
     
-    pipe = FluxImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-    pipe.to("cuda")
+    # If needed, add use_auth_token="YOUR_TOKEN" in from_pretrained below.
+    pipe = FluxImg2ImgPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16
+    ).to("cuda")
 
     generator = torch.Generator("cuda").manual_seed(seed)
     print(prompt)
-    output = pipe(prompt=prompt, image=image, generator=generator, strength=strength,
-                  guidance_scale=0, num_inference_steps=num_inference_steps, max_sequence_length=256)
+    output = pipe(
+        prompt=prompt,
+        image=image,
+        generator=generator,
+        strength=strength,
+        guidance_scale=0,
+        num_inference_steps=num_inference_steps,
+        max_sequence_length=256
+    )
 
-    # TODO: support mask
+    # TODO: Add mask support if needed
     return output.images[0]
 
 if __name__ == "__main__":
-    # args: input-image input-mask output
+    # Usage: python flux1_img2img.py input-image input-mask output
     image = Image.open(sys.argv[1])
     mask  = Image.open(sys.argv[2])
     output = process_image(image, mask)