Update app.py

app.py

@@ -1,12 +1,15 @@
+import spaces
 import gradio as gr
 import re
 from PIL import Image
+import os
+import numpy as np
 import torch
 from diffusers import FluxImg2ImgPipeline
 
-# Set up the device and pipeline
 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)
 
 def sanitize_prompt(prompt):
@@ -18,87 +21,100 @@ def convert_to_fit_size(original_width_and_height, maximum_size=2048):
     width, height = original_width_and_height
     if width <= maximum_size and height <= maximum_size:
         return width, height
-    
     scaling_factor = maximum_size / max(width, height)
-    return int(width * scaling_factor), int(height * scaling_factor)
+    new_width = int(width * scaling_factor)
+    new_height = int(height * scaling_factor)
+    return new_width, new_height
 
-def adjust_to_multiple_of_32(width, height):
-    return width - (width % 32), height - (height % 32)
+def adjust_to_multiple_of_32(width: int, height: int):
+    width = width - (width % 32)
+    height = height - (height % 32)
+    return width, height
 
+@spaces.GPU(duration=120)
 def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=4, progress=gr.Progress(track_tqdm=True)):
-    def process_img2img(image, prompt, strength, seed, num_inference_steps):
+    progress(0, desc="Starting")
+    def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
         if image is None:
             return None
         generator = torch.Generator(device).manual_seed(seed)
-        width, height = adjust_to_multiple_of_32(*convert_to_fit_size(image.size))
+        width, height = convert_to_fit_size(image.size)
+        width, height = adjust_to_multiple_of_32(width, 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
-        )
+        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)
         return output.images[0]
-    
-    return process_img2img(image, prompt, strength, seed, inference_step)
+    output = process_img2img(image, prompt, strength, seed, inference_step)
+    return output
+
+def read_file(path: str) -> str:
+    with open(path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    return content
 
-# Minimal CSS for black outline and container styling
 css = """
 #demo-container {
-    border: 2px solid black;
-    padding: 10px;
-    width: 100%;
-    max-width: 750px;
-    margin: auto;
+    border: 4px solid black;
+    border-radius: 8px;
+    padding: 20px;
+    margin: 20px auto;
+    max-width: 800px;
+}
+
+#image_upload, #output-img {
+    border: 4px solid black;
+    border-radius: 8px;
+    width: 256px;
+    height: 256px;
+    object-fit: cover;
 }
-#image_upload, #output-img, #generate_button {
-    border: 2px solid black;
+
+#run_button {
+    font-weight: bold;
+    border: 4px solid black;
+    border-radius: 8px;
+    padding: 10px 20px;
+}
+
+#col-left, #col-right {
+    max-width: 640px;
+    margin: 0 auto;
+}
+.grid-container {
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    gap: 10px;
+}
+.text {
+    font-size: 16px;
 }
 """
 
 with gr.Blocks(css=css, elem_id="demo-container") as demo:
     with gr.Column():
-        gr.HTML("<h1 style='text-align:center;'>Image to Image Generation</h1>")
+        gr.HTML(read_file("demo_header.html"))
+        # Removed or commented out the demo_tools.html line
+        # gr.HTML(read_file("demo_tools.html"))
     with gr.Row():
         with gr.Column():
-            image = gr.Image(
-                height=400, 
-                sources=['upload', 'clipboard'], 
-                image_mode='RGB', 
-                elem_id="image_upload", 
-                type="pil", 
-                label="Upload Image"
-            )
-            prompt = gr.Textbox(
-                label="Prompt", 
-                value="A woman", 
-                placeholder="Describe the output image", 
-                elem_id="prompt"
-            )
-            btn = gr.Button("Generate", elem_id="generate_button", variant="primary")
+            image = gr.Image(width=256, height=256, sources=['upload', 'clipboard'], image_mode='RGB', elem_id="image_upload", type="pil", label="Upload")
+            prompt = gr.Textbox(label="Prompt", value="a woman", placeholder="Your prompt", elem_id="prompt")
+            btn = gr.Button("Generate", elem_id="run_button", variant="primary")
             with gr.Accordion(label="Advanced Settings", open=False):
-                strength = gr.Number(value=0.75, minimum=0, maximum=1, step=0.01, label="Strength")
+                strength = gr.Number(value=0.75, minimum=0, maximum=0.75, step=0.01, label="Strength")
                 seed = gr.Number(value=100, minimum=0, step=1, label="Seed")
-                inference_step = gr.Number(value=4, minimum=1, step=1, label="Inference Steps")
+                inference_step = gr.Number(value=4, minimum=1, step=4, label="Inference Steps")
         with gr.Column():
-            image_out = gr.Image(
-                height=400, 
-                sources=[], 
-                label="Generated Output", 
-                elem_id="output-img", 
-                format="jpg"
-            )
+            image_out = gr.Image(width=256, height=256, label="Output", elem_id="output-img", format="jpg")
+
+    gr.HTML(gr.HTML(read_file("demo_footer.html")))
     
-    btn.click(
-        process_images, 
-        inputs=[image, prompt, strength, seed, inference_step], 
+    gr.on(
+        triggers=[btn.click, prompt.submit],
+        fn=process_images,
+        inputs=[image, prompt, strength, seed, inference_step],
         outputs=[image_out]
     )
 
-# Enable queue mode and CORS support
-demo.queue(concurrency_count=3, cors_allow_origins=["*"])
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()