Update app.py

app.py

@@ -1,8 +1,3 @@
-# Configuration
-prod = False
-port = 8080
-show_options = True  # Changed to True for better visibility
-
 import os
 import random
 import time
@@ -56,19 +51,17 @@ 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}")
-                     pass
+                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}")
         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)
 
@@ -173,7 +166,6 @@ def get_prompt(prompt, additional_prompt):
     return ", ".join(filter(None, prompt_parts))
 
 
-# Enhanced style list with more diverse options
 style_list = [
     {"name": "None",
      "prompt": ""
@@ -223,7 +215,6 @@ 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"
      },
@@ -247,7 +238,8 @@ 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())
 
@@ -256,7 +248,6 @@ def apply_style(style_name):
     return styles.get(style_name, "")
 
 
-# Enhanced CSS for Gradio UI
 css = """
     /* Global Styles */
     :root {
@@ -445,7 +436,6 @@ css = """
     }
     """
 
-# Load example images
 def load_examples():
     examples = []
     for i in range(1, 5):
@@ -461,13 +451,11 @@ 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>")
@@ -482,22 +470,16 @@ 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  # Set default image to in1.jpg
+                value=example_images[0] if example_images else None
             )
             
-            # 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)
-                        # Add click event for each example button
-                        btn.click(
-                            fn=lambda idx=i: select_example(idx), 
-                            outputs=image
-                        )
+                        btn.click(fn=lambda idx=i: select_example(idx), outputs=image)
             
         with gr.Column(scale=1, min_width=300):
             result = gr.Image(
@@ -508,7 +490,6 @@ 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(
@@ -520,7 +501,6 @@ 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"):
@@ -556,7 +536,6 @@ 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):
@@ -626,7 +605,17 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             elem_classes="helper-text"
         )
 
-    # Function to handle style selection changes across tabs
+    # 화면 맨 아래에 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>
+        """
+    )
+
     def update_style_selection(modern_value, classic_value, global_value):
         if modern_value is not None:
             return modern_value
@@ -637,7 +626,6 @@ 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):
@@ -649,7 +637,6 @@ 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],
@@ -668,7 +655,6 @@ 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
@@ -678,13 +664,11 @@ 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,
@@ -703,7 +687,6 @@ 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(),
@@ -715,18 +698,20 @@ 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
         )
 
@@ -741,27 +726,24 @@ 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,
@@ -771,7 +753,6 @@ 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, 
@@ -781,49 +762,24 @@ 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(
@@ -832,10 +788,13 @@ 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)
@@ -858,7 +817,6 @@ def process_image(
             image=control_image,
         ).images[0]
 
-    # Save and upload results (optional)
     try:
         timestamp = int(time.time())
         results_path = f"{timestamp}_output.jpg"
@@ -886,8 +844,9 @@ def process_image(
     return initial_result
 
 
-# Launch the Gradio app
+prod = False
+port = 8080
 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)