Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageEnhance
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -19,7 +19,7 @@ else:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt_part1, color, dress_type, front_design, back_design, prompt_part5, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def infer(prompt_part1, color, dress_type, front_design, back_design, prompt_part5, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     front_prompt = f"front view of {prompt_part1} {color} colored plain {dress_type} with {front_design} design, {prompt_part5}"
     back_prompt = f"back view of {prompt_part1} {color} colored plain {dress_type} with {back_design} design, {prompt_part5}"
     
@@ -30,7 +30,6 @@ def infer(prompt_part1, color, dress_type, front_design, back_design, prompt_par
     
     front_image = pipe(
         prompt=front_prompt, 
-        negative_prompt=negative_prompt,
         guidance_scale=guidance_scale, 
         num_inference_steps=num_inference_steps, 
         width=width, 
@@ -40,7 +39,6 @@ def infer(prompt_part1, color, dress_type, front_design, back_design, prompt_par
     
     back_image = pipe(
         prompt=back_prompt, 
-        negative_prompt=negative_prompt,
         guidance_scale=guidance_scale, 
         num_inference_steps=num_inference_steps, 
         width=width, 
@@ -50,17 +48,6 @@ def infer(prompt_part1, color, dress_type, front_design, back_design, prompt_par
     
     return front_image, back_image
 
-examples = [
-    ["red", "t-shirt", "yellow stripes", "polka dots"],
-    ["blue", "hoodie", "minimalist", "abstract art"],
-    ["red", "sweat shirt", "geometric design", "plain"],
-]
-
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 def edit_image(img_data, operation, *args):
     image = Image.open(img_data)
     
@@ -78,48 +65,68 @@ def edit_image(img_data, operation, *args):
             image = ImageOps.mirror(image)
         else:
             image = ImageOps.flip(image)
+    elif operation == "color":
+        factor = float(args[0])
+        image = ImageEnhance.Color(image).enhance(factor)
     
     return image
 
+examples = [
+    ["red", "t-shirt", "yellow stripes", "polka dots"],
+    ["blue", "hoodie", "minimalist", "abstract art"],
+    ["red", "sweat shirt", "geometric design", "plain"],
+]
+
+if torch.cuda.is_available():
+    power_device = "GPU"
+else:
+    power_device = "CPU"
+
 with gr.Blocks() as demo:
-    
-    with gr.Column():
+    with gr.Row():
         gr.Markdown(f"""
         # GenZ Couture
         Currently running on {power_device}.
         """)
-        
-        prompt_part1 = gr.Textbox(value="a single", label="Prompt Part 1", visible=False)
-        prompt_part2 = gr.Textbox(label="color", placeholder="color (e.g., red, blue)")
-        prompt_part3 = gr.Textbox(label="dress_type", placeholder="dress_type (e.g., t-shirt, hoodie)")
-        prompt_part4_front = gr.Textbox(label="front design", placeholder="front design")
-        prompt_part4_back = gr.Textbox(label="back design", placeholder="back design")
-        prompt_part5 = gr.Textbox(value="hanging on the plain wall", label="Prompt Part 5", visible=False)
-        
-        run_button = gr.Button("Generate Designs")
 
-        front_result = gr.Image(label="Front View Result", type="pil", interactive=True)
-        back_result = gr.Image(label="Back View Result", type="pil", interactive=True)
-
-        gr.Examples(examples=examples, inputs=[prompt_part2, prompt_part3, prompt_part4_front, prompt_part4_back])
+    with gr.Row():
+        with gr.Column():
+            prompt_part2 = gr.Textbox(label="Color", placeholder="Color (e.g., red, blue)")
+            prompt_part3 = gr.Textbox(label="Dress Type", placeholder="Dress Type (e.g., t-shirt, hoodie)")
+            prompt_part4_front = gr.Textbox(label="Front Design", placeholder="Front Design")
+            prompt_part4_back = gr.Textbox(label="Back Design", placeholder="Back Design")
+            seed = gr.Slider(0, MAX_SEED, step=1, label="Seed", value=42)
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+            width = gr.Slider(256, MAX_IMAGE_SIZE, step=32, label="Width", value=512)
+            height = gr.Slider(256, MAX_IMAGE_SIZE, step=32, label="Height", value=512)
+            guidance_scale = gr.Slider(1, 20, step=0.5, label="Guidance Scale", value=7.5)
+            num_inference_steps = gr.Slider(10, 100, step=1, label="Number of Inference Steps", value=50)
+            
+            run_button = gr.Button("Generate Designs")
+        
+        with gr.Column():
+            front_result = gr.Image(label="Front View Result", type="pil", interactive=True)
+            back_result = gr.Image(label="Back View Result", type="pil", interactive=True)
 
     run_button.click(
         fn=infer,
-        inputs=[prompt_part1, prompt_part2, prompt_part3, prompt_part4_front, prompt_part4_back, prompt_part5],
+        inputs=["a single", prompt_part2, prompt_part3, prompt_part4_front, prompt_part4_back, "hanging on the plain wall", seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs=[front_result, back_result]
     )
 
     gr.Markdown("## Creative Touch")
 
-    edit_operation = gr.Dropdown(choices=["rotate", "crop", "resize", "flip"], label="Edit Operation")
-    edit_args = gr.Textbox(label="Edit Arguments (comma-separated)", placeholder="For rotate: angle, For crop: left,top,right,bottom, For resize: width,height, For flip: horizontal/vertical")
-    
-    edit_button = gr.Button("Edit Front Design")
+    with gr.Row():
+        edit_operation = gr.Dropdown(choices=["rotate", "crop", "resize", "flip", "color"], label="Edit Operation")
+        edit_args = gr.Textbox(label="Edit Arguments (comma-separated)", placeholder="For rotate: angle, For crop: left,top,right,bottom, For resize: width,height, For flip: horizontal/vertical, For color: factor")
 
-    edit_button.click(
-        fn=lambda img_data, operation, args: edit_image(img_data, operation, *args.split(',')),
-        inputs=[front_result, edit_operation, edit_args],
-        outputs=[front_result]
-    )
+        edit_button = gr.Button("Edit Front Design")
+        edited_image = gr.Image(label="Edited Front Design", type="pil", interactive=True)
+
+        edit_button.click(
+            fn=lambda img_data, operation, args: edit_image(img_data, operation, *args.split(',')),
+            inputs=[front_result, edit_operation, edit_args],
+            outputs=[edited_image]
+        )
 
 demo.queue().launch()