Update app.py

app.py

@@ -1,3 +1,8 @@
+# Configuration
+prod = False
+port = 8080
+show_options = True  # Changed to True for better visibility
+
 import os
 import random
 import time
@@ -51,17 +56,19 @@ class Preprocessor:
         
         if hasattr(self.model, 'device'):
             if self.model.device.type != device:
-                print(f"Moving preprocessor model to {device}")
-                try:
-                    self.model.to(device)
-                except Exception as e:
-                    print(f"Error moving preprocessor model to {device}: {e}")
+                 print(f"Moving preprocessor model to {device}")
+                 try:
+                     self.model.to(device)
+                 except Exception as e:
+                     print(f"Error moving preprocessor model to {device}: {e}")
+                     pass
         else:
             print("Warning: Preprocessor model has no .device attribute. Attempting to move to correct device.")
             try:
                 self.model.to(device)
             except Exception as e:
                 print(f"Error attempting to move preprocessor model without .device attribute: {e}")
+                pass
 
         return self.model(image, **kwargs)
 
@@ -166,6 +173,7 @@ def get_prompt(prompt, additional_prompt):
     return ", ".join(filter(None, prompt_parts))
 
 
+# Enhanced style list with more diverse options
 style_list = [
     {"name": "None",
      "prompt": ""
@@ -215,6 +223,7 @@ style_list = [
     {"name": "Matrix",
      "prompt": "Futuristic cyberpunk interior,neon accent lighting,holographic plants,sleek black surfaces,advanced gaming setup,transparent screens,Blade Runner inspired decor,high-tech minimalist furniture"
      },
+    # New added styles
     {"name": "Industrial Loft",
      "prompt": "Industrial loft interior,exposed brick walls,metal finishes,high ceilings with exposed pipes,concrete floors,vintage factory lights,open floor plan"
      },
@@ -238,8 +247,7 @@ style_list = [
      },
     {"name": "Penthouse",
      "prompt": "Luxury penthouse interior,floor-to-ceiling windows,city skyline views,modern furniture,high-end appliances,marble countertops,designer lighting fixtures"
-     }
-]
+     }]
 styles = {k["name"]: (k["prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
 
@@ -248,6 +256,7 @@ def apply_style(style_name):
     return styles.get(style_name, "")
 
 
+# Enhanced CSS for Gradio UI
 css = """
     /* Global Styles */
     :root {
@@ -436,6 +445,7 @@ css = """
     }
     """
 
+# Load example images
 def load_examples():
     examples = []
     for i in range(1, 5):
@@ -451,11 +461,13 @@ def load_examples():
 
 example_images = load_examples()
 
+# Function to select example image
 def select_example(index):
     if 0 <= index < len(example_images):
         return example_images[index]
     return None
 
+# Gradio Interface Definition
 with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     gr.Markdown("<h1>✨ Dream of IKEA ✨</h1>")
     gr.Markdown("<h3>Transform your space with AI-powered interior design</h3>")
@@ -470,16 +482,22 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 mirror_webcam=True,
                 type="pil",
                 elem_id="input-image",
-                value=example_images[0] if example_images else None
+                value=example_images[0] if example_images else None  # Set default image to in1.jpg
             )
             
+            # Example images section with buttons instead of images with style
             with gr.Row(elem_classes="example-images"):
+                # Create example buttons
                 example_buttons = []
                 for i in range(len(example_images)):
                     if example_images[i]:
                         btn = gr.Button(f"Example {i+1}", elem_classes="example-thumb")
                         example_buttons.append(btn)
-                        btn.click(fn=lambda idx=i: select_example(idx), outputs=image)
+                        # Add click event for each example button
+                        btn.click(
+                            fn=lambda idx=i: select_example(idx), 
+                            outputs=image
+                        )
             
         with gr.Column(scale=1, min_width=300):
             result = gr.Image(
@@ -490,6 +508,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 elem_id="output-image"
             )
     
+    # Design input section
     with gr.Row():
         with gr.Column(scale=2):
             prompt = gr.Textbox(
@@ -501,6 +520,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             run_button = gr.Button(value="🚀 Generate Design", size="lg")
             use_ai_button = gr.Button(value="♻️ Use Result as New Input", size="lg")
     
+    # Grouped style selection with categories
     gr.Markdown("<h2>Design Style Selection</h2>")
     with gr.Tabs():
         with gr.TabItem("Modern Styles"):
@@ -536,6 +556,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 elem_id="global-styles"
             )
     
+    # Advanced options - now with a clearer separator and improved layout
     with gr.Accordion("⚙️ Advanced Options", open=False):
         with gr.Row():
             with gr.Column(scale=1):
@@ -605,17 +626,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             elem_classes="helper-text"
         )
 
-    # 화면 맨 아래에 Discord 배지를 삽입
-    gr.Markdown(
-        """
-        <p style="text-align: center;">
-            <a href="https://discord.gg/openfreeai" target="_blank">
-                <img src="https://img.shields.io/static/v1?label=Discord&message=Join%20our%20community&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="badge">
-            </a>
-        </p>
-        """
-    )
-
+    # Function to handle style selection changes across tabs
     def update_style_selection(modern_value, classic_value, global_value):
         if modern_value is not None:
             return modern_value
@@ -626,6 +637,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         else:
             return "None"
     
+    # Style synchronization
     style_selection = gr.State("None")
     
     def clear_other_tabs(active_tab, value):
@@ -637,6 +649,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             return None, None, value
         return None, None, None
     
+    # Connect the tab radios to update each other
     style_selection_modern.change(
         fn=lambda x: clear_other_tabs("modern", x),
         inputs=[style_selection_modern],
@@ -655,6 +668,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         outputs=[style_selection_modern, style_selection_classic, style_selection_global]
     )
     
+    # Combine all style selections into one for processing
     def get_active_style(modern, classic, global_style):
         if modern is not None and modern != "":
             return modern
@@ -664,11 +678,13 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             return global_style
         return "None"
     
+    # Randomize seed function
     def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         return seed
 
+    # Configuration list for inputs - using function to get active style
     def get_config_inputs():
         return [
             image,
@@ -687,6 +703,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             randomize_seed,
         ]
 
+    # Gradio Event Handling Functions
     @gr.on(
         triggers=[image.upload, prompt.submit, run_button.click], 
         inputs=get_config_inputs(),
@@ -698,20 +715,18 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale,
         seed, randomize_seed, progress=gr.Progress(track_tqdm=True)
     ):
+        # Get the active style
         active_style = get_active_style(style_modern, style_classic, style_global)
+        
+        # Apply seed randomization
         processed_seed = randomize_seed_fn(seed, randomize_seed)
         print(f"Using processed seed: {processed_seed}")
+        print(f"Active style: {active_style}")
+
+        # Call the core processing function
         return process_image(
-            image, 
-            active_style, 
-            prompt, 
-            a_prompt, 
-            n_prompt, 
-            num_images,
-            image_resolution, 
-            preprocess_resolution, 
-            num_steps, 
-            guidance_scale,
+            image, active_style, prompt, a_prompt, n_prompt, num_images,
+            image_resolution, preprocess_resolution, num_steps, guidance_scale,
             processed_seed
         )
 
@@ -726,24 +741,27 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, 
         num_steps, guidance_scale, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)
     ):
+        # First, yield the previous result to update the input image immediately
         yield previous_result, gr.update()
+
+        # Get active style
         active_style = get_active_style(style_modern, style_classic, style_global)
+        
+        # Apply seed randomization
         processed_seed = randomize_seed_fn(seed, randomize_seed)
+        print(f"Using processed seed: {processed_seed}")
+
+        # Then, process the new input image
         new_result = process_image(
-            previous_result, 
-            active_style, 
-            prompt, 
-            a_prompt,
-            n_prompt, 
-            num_images, 
-            image_resolution,
-            preprocess_resolution, 
-            num_steps, 
-            guidance_scale,
+            previous_result, active_style, prompt, a_prompt,
+            n_prompt, num_images, image_resolution,
+            preprocess_resolution, num_steps, guidance_scale,
             processed_seed
         )
+        # Finally, yield the new result
         yield previous_result, new_result
 
+    # Turn off buttons when processing
     @gr.on(
         triggers=[image.upload, use_ai_button.click, run_button.click], 
         inputs=None,
@@ -753,6 +771,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     def turn_buttons_off():
         return gr.update(interactive=False, value="Processing..."), gr.update(interactive=False)
 
+    # Turn on buttons when processing is complete
     @gr.on(
         triggers=[result.change], 
         inputs=None, 
@@ -762,24 +781,49 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     def turn_buttons_on():
         return gr.update(interactive=True, value="♻️ Use Result as New Input"), gr.update(interactive=True, value="🚀 Generate Design")
 
+
+# Core Image Processing Function
+@spaces.GPU(duration=12)
+@torch.inference_mode()
 def process_image(
-    image,
-    style_selection,
-    prompt,
-    a_prompt,
-    n_prompt,
-    num_images,
-    image_resolution,
-    preprocess_resolution,
-    num_steps,
-    guidance_scale,
-    seed,
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        guidance_scale,
+        seed,
 ):
+    """
+    Processes an input image to generate a new image based on style and prompts.
+
+    Args:
+        image: Input PIL Image.
+        style_selection: Name of the design style to apply.
+        prompt: Custom design prompt.
+        a_prompt: Additional positive prompt.
+        n_prompt: Negative prompt.
+        num_images: Number of images to generate (currently only 1 supported by pipeline).
+        image_resolution: Resolution for the output image.
+        preprocess_resolution: Resolution for the preprocessor.
+        num_steps: Number of inference steps.
+        guidance_scale: Guidance scale for the diffusion process.
+        seed: Random seed for reproducibility.
+
+    Returns:
+        A PIL Image of the generated result.
+    """
+    # Use the seed passed from the event handler
     current_seed = seed
     generator = torch.cuda.manual_seed(current_seed) if torch.cuda.is_available() else torch.manual_seed(current_seed)
 
     if preprocessor.name != "NormalBae":
-        preprocessor.load("NormalBae")
+         preprocessor.load("NormalBae")
+
     preprocessor.model.to("cuda" if torch.cuda.is_available() else "cpu")
 
     control_image = preprocessor(
@@ -788,13 +832,10 @@ def process_image(
         detect_resolution=preprocess_resolution,
     )
 
+    # Construct the full prompt
     if style_selection and style_selection != "None":
         style_prompt = apply_style(style_selection)
-        prompt_parts = [
-            f"Photo from Pinterest of {prompt}" if prompt else None, 
-            style_prompt if style_prompt else None, 
-            a_prompt if a_prompt else None
-        ]
+        prompt_parts = [f"Photo from Pinterest of {prompt}" if prompt else None, style_prompt if style_prompt else None, a_prompt if a_prompt else None]
         full_prompt = ", ".join(filter(None, prompt_parts))
     else:
         full_prompt = get_prompt(prompt, a_prompt)
@@ -817,6 +858,7 @@ def process_image(
             image=control_image,
         ).images[0]
 
+    # Save and upload results (optional)
     try:
         timestamp = int(time.time())
         results_path = f"{timestamp}_output.jpg"
@@ -844,9 +886,8 @@ def process_image(
     return initial_result
 
 
-prod = False
-port = 8080
+# Launch the Gradio app
 if prod:
     demo.queue(max_size=20).launch(server_name="localhost", server_port=port)
 else:
-    demo.queue().launch(share=True, show_api=False)
+    demo.queue().launch(share=True, show_api=False)