Update app.py

app.py

@@ -8,10 +8,15 @@ import numpy as np
 import torch
 from diffusers import FluxImg2ImgPipeline
 
+# Set the torch data type and device
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(device)
+# If needed, add use_auth_token="YOUR_TOKEN" in from_pretrained below.
+pipe = FluxImg2ImgPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+).to(device)
 
 def sanitize_prompt(prompt):
     # Allow only alphanumeric characters, spaces, and basic punctuation
@@ -44,10 +49,10 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
 
     def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
         if image is None:
-            print("empty input image returned")
+            print("Empty input image returned.")
             return None
 
-        # Ensure image is in RGB mode (this helps with WebP and other formats)
+        # Ensure the image is in RGB mode (handles formats like WebP and JFIF)
         if image.mode != "RGB":
             image = image.convert("RGB")
         
@@ -56,8 +61,17 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
         width, height = adjust_to_multiple_of_32(fit_width, fit_height)
         image = image.resize((width, height), Image.LANCZOS)
     
-        output = pipe(prompt=prompt, image=image, generator=generator, strength=strength, width=width, height=height,
-                      guidance_scale=0, num_inference_steps=num_inference_steps, max_sequence_length=256)
+        output = pipe(
+            prompt=prompt,
+            image=image,
+            generator=generator,
+            strength=strength,
+            width=width,
+            height=height,
+            guidance_scale=0,
+            num_inference_steps=num_inference_steps,
+            max_sequence_length=256,
+        )
     
         pil_image = output.images[0]
         new_width, new_height = pil_image.size
@@ -66,7 +80,7 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
             resized_image = pil_image.resize((fit_width, fit_height), Image.LANCZOS)
             return resized_image
         return pil_image
-    
+
     output = process_img2img(image, prompt, strength, seed, inference_step)
     return output
 
@@ -88,7 +102,7 @@ css = """
   display: flex;
   align-items: center;
   justify-content: center;
-  gap:10px
+  gap: 10px;
 }
 .image {
   width: 128px; 
@@ -106,10 +120,22 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
         gr.HTML(read_file("demo_tools.html"))
     with gr.Row():
         with gr.Column():
-            image = gr.Image(height=800, sources=['upload','clipboard'], image_mode='RGB', elem_id="image_upload", type="pil", label="Upload")
+            image = gr.Image(
+                height=800,
+                sources=['upload', 'clipboard'],
+                image_mode='RGB',
+                elem_id="image_upload",
+                type="pil",
+                label="Upload"
+            )
             with gr.Row(elem_id="prompt-container", equal_height=False):
                 with gr.Row():
-                    prompt = gr.Textbox(label="Prompt", value="a women", placeholder="Your prompt (what you want in place of what is erased)", elem_id="prompt")
+                    prompt = gr.Textbox(
+                        label="Prompt",
+                        value="a women",
+                        placeholder="Your prompt (what you want in place of what is erased)",
+                        elem_id="prompt"
+                    )
             btn = gr.Button("Img2Img", elem_id="run_button", variant="primary")
             with gr.Accordion(label="Advanced Settings", open=False):
                 with gr.Row(equal_height=True):