Update app.py

app.py

@@ -2,76 +2,59 @@ import spaces
 import gradio as gr
 import re
 from PIL import Image
-
 import os
 import numpy as np
 import torch
-from diffusers import FluxImg2ImgPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
 
-# Use default float32 precision for CPU
-dtype = torch.float32
-device = "cpu"
+# Use float16 for lower VRAM usage
+dtype = torch.float16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the pipeline on CPU
-pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+# Load the lighter model on the GPU
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=dtype)
+pipe.to(device)
 
 def sanitize_prompt(prompt):
-    # Allow only alphanumeric characters, spaces, and basic punctuation
     allowed_chars = re.compile(r"[^a-zA-Z0-9\s.,!?-]")
-    sanitized_prompt = allowed_chars.sub("", prompt)
-    return sanitized_prompt
+    return allowed_chars.sub("", prompt)
 
 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
-    
-    if width > height:
-        scaling_factor = maximum_size / width
-    else:
-        scaling_factor = maximum_size / height
-
-    new_width = int(width * scaling_factor)
-    new_height = int(height * scaling_factor)
-    return new_width, new_height
+    scaling_factor = maximum_size / max(width, height)
+    return int(width * scaling_factor), int(height * scaling_factor)
 
 def adjust_to_multiple_of_32(width: int, height: int):
-    width = width - (width % 32)
-    height = height - (height % 32)
-    return width, height
-
-@spaces.CPU(duration=120)
-def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=4, progress=gr.Progress(track_tqdm=True)):
-    progress(0, desc="Starting")
+    return width - (width % 32), height - (height % 32)
 
-    def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
+@spaces.GPU(duration=120)
+def process_images(image, prompt="a woman", strength=0.75, seed=0, inference_step=50, progress=gr.Progress(track_tqdm=True)):
+    progress(0, desc="Starting processing")
+    
+    def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=50):
         if image is None:
-            print("empty input image returned")
+            print("Empty input image returned")
             return None
-        # Create a CPU generator
-        generator = torch.Generator("cpu").manual_seed(seed)
+        generator = torch.Generator(device).manual_seed(seed)
         fit_width, fit_height = convert_to_fit_size(image.size)
         width, height = adjust_to_multiple_of_32(fit_width, fit_height)
         image = image.resize((width, height), Image.LANCZOS)
     
         output = pipe(
             prompt=prompt,
-            image=image,
+            init_image=image,
             generator=generator,
             strength=strength,
-            width=width,
-            height=height,
-            guidance_scale=0,
+            guidance_scale=7.5,
             num_inference_steps=num_inference_steps,
-            max_sequence_length=256
         )
     
         pil_image = output.images[0]
-        new_width, new_height = pil_image.size
-    
-        if (new_width != fit_width) or (new_height != fit_height):
-            resized_image = pil_image.resize((fit_width, fit_height), Image.LANCZOS)
-            return resized_image
+        # Optionally, resize back to original fitted dimensions if desired
+        if pil_image.size != (fit_width, fit_height):
+            pil_image = pil_image.resize((fit_width, fit_height), Image.LANCZOS)
         return pil_image
     
     output = process_img2img(image, prompt, strength, seed, inference_step)
@@ -79,8 +62,7 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
 
 def read_file(path: str) -> str:
     with open(path, 'r', encoding='utf-8') as f:
-        content = f.read()
-    return content
+        return f.read()
 
 css = """
 #col-left {
@@ -95,15 +77,13 @@ css = """
   display: flex;
   align-items: center;
   justify-content: center;
-  gap:10px
+  gap: 10px;
 }
-
 .image {
   width: 128px; 
   height: 128px; 
   object-fit: cover; 
 }
-
 .text {
   font-size: 16px;
 }
@@ -115,40 +95,26 @@ 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 woman", placeholder="Your prompt", 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):
-                    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=4, label="inference_step")
+                    strength = gr.Number(value=0.75, minimum=0, maximum=1.0, step=0.01, label="Strength")
+                    seed = gr.Number(value=100, minimum=0, step=1, label="Seed")
+                    inference_step = gr.Number(value=50, minimum=1, step=1, label="Inference Steps")
                 id_input = gr.Text(label="Name", visible=False)
-                            
         with gr.Column():
             image_out = gr.Image(height=800, sources=[], label="Output", elem_id="output-img", format="jpg")
     
     gr.Examples(
         examples=[
-            ["examples/draw_input.jpg", "examples/draw_output.jpg", "a women ,eyes closed,mouth opened"],
-            ["examples/draw-gimp_input.jpg", "examples/draw-gimp_output.jpg", "a women ,eyes closed,mouth opened"],
-            ["examples/gimp_input.jpg", "examples/gimp_output.jpg", "a women ,hand on neck"],
-            ["examples/inpaint_input.jpg", "examples/inpaint_output.jpg", "a women ,hand on neck"]
+            ["examples/draw_input.jpg", "examples/draw_output.jpg", "a woman with blue eyes"],
+            ["examples/draw-gimp_input.jpg", "examples/draw-gimp_output.jpg", "a woman with a serene expression"],
+            ["examples/gimp_input.jpg", "examples/gimp_output.jpg", "a woman in a garden"],
+            ["examples/inpaint_input.jpg", "examples/inpaint_output.jpg", "a woman in a futuristic city"]
         ],
         inputs=[image, image_out, prompt],
     )