Update flux1_img2img.py

flux1_img2img.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import FluxImg2ImgPipeline
 from PIL import Image
 import sys
 import spaces
@@ -7,43 +7,59 @@ import spaces
 @spaces.GPU
 def process_image(
     image,
-    mask_image,  # unused, but we keep it in signature
+    mask_image,
     prompt="a person",
-    model_id="SG161222/Realistic_Vision_V2.0",
+    model_id="black-forest-labs/FLUX.1-schnell",
     strength=0.75,
     seed=0,
-    num_inference_steps=30
+    num_inference_steps=4
 ):
+    print("start process image process_image")
     if image is None:
-        print("No input image provided.")
+        print("empty input image returned")
         return None
 
-    # Load Realistic Vision model from Hugging Face
-    # Typically doesn't require "Accept" step
-    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-        model_id, torch_dtype=torch.float16
+    # 1) Use float16 (T4 doesn't have native bf16 support)
+    # 2) low_cpu_mem_usage=True for more efficient loading
+    # 3) Optionally enable xFormers
+    pipe = FluxImg2ImgPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        revision="fp16",            # sometimes needed if the repo has an FP16 branch
+        low_cpu_mem_usage=True
     )
+
+    # Move to GPU
     pipe.to("cuda")
 
-    # Set a seed for reproducibility
+    # If you have xFormers installed (pip install xformers):
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+        print("Enabled xFormers memory efficient attention.")
+    except Exception as e:
+        print("xFormers not available:", e)
+
+    # Create a reproducible generator
     generator = torch.Generator("cuda").manual_seed(seed)
 
-    # Increase guidance_scale to ~7–10 for more prompt adherence
+    print(f"Prompt: {prompt}")
     output = pipe(
         prompt=prompt,
         image=image,
+        generator=generator,
         strength=strength,
-        guidance_scale=8.5,
+        guidance_scale=0,  # same as your original code
         num_inference_steps=num_inference_steps,
-        generator=generator
+        max_sequence_length=256
     )
 
+    # TODO: support mask if needed
     return output.images[0]
 
 if __name__ == "__main__":
-    # Usage: python process_image.py <input_image> <mask_image> <output_image>
+    # Usage: python img2img.py input_image.png input_mask.png output.png
     image = Image.open(sys.argv[1]).convert("RGB")
-    mask = Image.open(sys.argv[2]).convert("RGB")  # Not used in this example
+    mask = Image.open(sys.argv[2]).convert("RGB")  # currently unused
     result = process_image(image, mask)
     if result is not None:
         result.save(sys.argv[3])