Update app.py

app.py

@@ -1,11 +1,15 @@
 # Configuration
-# These variables are now mostly for reference, FastAPI/Uvicorn handle port
 prod = False
-import os # Moved import os to the beginning
-port = int(os.environ.get("PORT", 8080)) # Use PORT environment variable provided by Spaces, default to 8080
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
 
+import os
 import random
 import time
+import gradio as gr # Re-added gradio import
 import numpy as np
 import spaces
 import imageio
@@ -18,16 +22,10 @@ from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
+    # StableDiffusionInpaintPipeline, # Commented out as inpainting part was commented
+    # AutoencoderKL, # Commented out as VAE part was commented
 )
-# Assuming controlnet_aux_local is a local package or needs to be installed separately
-from controlnet_aux_local import NormalBaeDetector
-
-# Import necessary components for FastAPI
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
-import uvicorn
-import io
+from controlnet_aux_local import NormalBaeDetector # Assuming this local package is available
 
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
@@ -37,8 +35,6 @@ print("loading everything")
 compiled = False
 api = HfApi()
 
-# Initialize FastAPI app
-app = FastAPI()
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
@@ -69,7 +65,6 @@ class Preprocessor:
 
 # Load models and preprocessor when the script starts
 # Controlnet Normal
-# Corrected typo in the model ID
 model_id = "lllyasviel/control_v11p_sd15_normalbae"
 print("initializing controlnet")
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -162,21 +157,13 @@ def get_prompt(prompt, additional_prompt):
     if prompt:
         prompt_parts.append(f"Photo from Pinterest of {prompt}")
     else:
-        # If no specific prompt, use a default or random one (original code's 'boho chic' or random 'girls' prompts)
-        # Let's stick to interior design context, so maybe a default interior style if no prompt?
-        # Or, based on the original code's `if prompt == "":` block, it seemed to sometimes
-        # default to random 'girl' prompts. This might be unintended for an interior design API.
-        # Let's assume if no prompt is given, we still apply the interior context.
-        prompt_parts.append("Photo from Pinterest of interior space") # Default if no prompt
+        # If no specific prompt, use a default interior context
+        prompt_parts.append("Photo from Pinterest of interior space")
 
     prompt_parts.append(interior)
 
     if additional_prompt:
         prompt_parts.append(additional_prompt)
-    # Note: The original `get_prompt` had a block that randomly selected 'girl' related prompts
-    # when the input `prompt` was empty. This seems out of place for an interior design API.
-    # I have removed that random selection logic to focus on interior design prompts.
-    # If you need that random girl prompt functionality, please clarify where/how it should be used.
 
     return ", ".join(filter(None, prompt_parts))
 
@@ -238,40 +225,238 @@ def apply_style(style_name):
     return styles.get(style_name, "")
 
 
-# The core processing function, now called by the API endpoint
-@torch.inference_mode() # Keep inference_mode here for efficiency
-def process_image_api(
-    image: Image.Image,
-    style_selection: str = "None",
-    prompt: str = "",
-    a_prompt: str = "",
-    n_prompt: str = "EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
-    image_resolution: int = 512,
-    preprocess_resolution: int = 512,
-    num_steps: int = 15,
-    guidance_scale: float = 5.5,
-    seed: int = -1,
-):
-    """
-    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.
-        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. Use -1 for random seed.
-
-    Returns:
-        A PIL Image of the generated result.
+# CSS for Gradio UI
+css = """h1, h2, h3 {
+        text-align: center;
+        display: block;
+    }
+
+    footer {
+        visibility: hidden;
+    }
+
+    .gradio-container {
+        max-width: 1100px !important;
+    }
+
+    .gr-image {
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        width: 100%;
+        height: 512px;
+        overflow: hidden;
+    }
+
+    .gr-image img {
+        width: 100%;
+        height: 100%;
+        object-fit: cover;
+        object-position: center;
+    }
     """
-    current_seed = seed if seed != -1 else random.randint(0, MAX_SEED)
-    generator = torch.cuda.manual_seed(current_seed) if torch.cuda.is_available() else torch.manual_seed(current_seed)
+
+# Gradio Interface Definition
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    gr.Markdown("## Interior AI Designer") # Added a title
+    #############################################################################
+    with gr.Row():
+        # Re-added Accordion with visible=show_options
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            a_prompt = gr.Textbox(
+                label="Additional prompt",
+                value="design-style interior designed (interior space), tungsten white balance, captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                value="EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
+            )
+    #############################################################################
+    # input text
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Design",
+            placeholder="Enter a description (optional)",
+        )
+    # design options
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1, min_width=300):
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                type="pil",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size="lg", visible=False)
+        # output image
+        with gr.Column(scale=1, min_width=300):
+            result = gr.Image(
+                label="Output",
+                interactive=False,
+                type="pil",
+                show_share_button=False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size="lg", visible=False)
+    # Configuration list for inputs
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        guidance_scale,
+        seed,
+    ]
+    with gr.Row():
+        helper_text = gr.Markdown("## Tap and hold (on mobile) to save the image.", visible=True)
+
+    # Gradio Event Handling Functions
+    @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config,
+           outputs=result, show_progress="minimal")
+    def auto_process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images,
+                           image_resolution, preprocess_resolution, num_steps, guidance_scale,
+                           seed, progress=gr.Progress(track_tqdm=True)):
+        # Call the core processing function
+        return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images,
+                             image_resolution, preprocess_resolution, num_steps, guidance_scale,
+                             seed)
+
+    @gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result],
+           show_progress="minimal")
+    def submit(previous_result, image, style_selection, prompt, a_prompt, n_prompt, num_images,
+               image_resolution, preprocess_resolution, num_steps, guidance_scale, seed,
+               progress=gr.Progress(track_tqdm=True)):
+        # First, yield the previous result to update the input image immediately
+        yield previous_result, gr.update()
+        # Then, process the new input image
+        new_result = process_image(previous_result, style_selection, prompt, a_prompt,
+                                   n_prompt, num_images, image_resolution,
+                                   preprocess_resolution, num_steps, guidance_scale, seed)
+        # Finally, yield the new result
+        yield previous_result, new_result
+
+    @gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None,
+           outputs=[run_button, use_ai_button], show_progress="hidden")
+    def turn_buttons_off():
+        return gr.update(visible=False), gr.update(visible=False)
+
+    @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button],
+           show_progress="hidden")
+    def turn_buttons_on():
+        return gr.update(visible=True), gr.update(visible=True)
+
+
+# Core Image Processing Function (renamed back from process_image_api)
+@spaces.GPU(duration=12) # Re-added spaces.GPU decorator
+@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,
+):
+    # Use provided seed or generate a random one based on randomize_seed checkbox (need to add randomize_seed input)
+    # The original Gradio code used randomize_seed_fn, but the logic is simple enough to keep here
+    # Let's add randomize_seed as an input parameter to this function
+    # Re-checking the original Gradio code, randomize_seed is an input to the Gradio event function,
+    # but not directly passed to process_image. Instead, the seed is randomized *before* calling process_image
+    # in the Gradio event handlers. Let's replicate that.
+
+    # Seed randomization logic from original Gradio code
+    # The seed parameter in the config list is used here
+    current_seed = seed
+    # Note: The original `auto_process_image` and `submit` functions did NOT use the `randomize_seed` checkbox
+    # to potentially override the seed slider value before passing it to `process_image`.
+    # The `randomize_seed_fn` was defined but not called in the provided Gradio code snippet.
+    # Let's stick to the provided Gradio code's logic which seems to just use the seed slider value.
+    # If randomization is needed, it should be handled in the event handler before calling process_image.
+    # However, the original `process_image` function *itself* had `seed = random.randint(0, MAX_SEED)`.
+    # This means the seed slider was ignored! Let's fix that and use the seed slider value,
+    # applying randomization if the checkbox is checked in the event handler.
+
+    # Re-adding randomize_seed to the inputs list `config` and event handler functions
+    # Then applying randomization logic in the event handler before calling process_image
+
+    # *** Correction: Re-reading the original Gradio code snippet provided at the very beginning,
+    # the `randomize_seed_fn` was defined but not used. The `process_image` function
+    # *inside* the `if gr.NO_RELOAD:` block *did* have `seed = random.randint(0, MAX_SEED)`
+    # at the beginning, effectively ignoring the input seed slider value unless randomize_seed_fn
+    # was somehow implicitly called or the logic was elsewhere.
+    # Let's assume the intention was to use the seed slider, with an option to randomize.
+    # The `randomize_seed_fn` should be called in the Gradio event handlers.
+
+    # Let's modify the Gradio event handlers to use randomize_seed_fn
+
+    # *** Another Correction: The `process_image` function signature in the original Gradio code
+    # included `seed` as a parameter, but the first line inside the function was `seed = random.randint(0, MAX_SEED)`.
+    # This is contradictory. It implies the input `seed` was always overwritten by a random value.
+    # Let's assume the *intent* was to use the input `seed` unless `randomize_seed` was True.
+    # The `randomize_seed_fn` should be called in the event handler to get the final seed value.
+
+    # Let's modify the event handlers to call randomize_seed_fn and pass the result as `seed` to `process_image`.
+    # The `process_image` function itself should *not* randomize the seed internally if a seed is passed.
+
+    # *** Final Decision: Let's revert `process_image` to its signature from the original Gradio code,
+    # but remove the internal `seed = random.randint(0, MAX_SEED)` line.
+    # The randomization logic will be added to the Gradio event handlers using `randomize_seed_fn`.
+    # The `config` list and event handlers will need `randomize_seed` as an input.
+
+    # Re-adding randomize_seed to config and event handler signatures
+    # Adding call to randomize_seed_fn in event handlers
+
+    # --- Start of process_image logic ---
+    # Seed handling moved to event handlers
 
     if preprocessor.name != "NormalBae":
          preprocessor.load("NormalBae")
@@ -296,18 +481,21 @@ def process_image_api(
     negative_prompt = str(n_prompt)
     print(f"Using prompt: {full_prompt}")
     print(f"Using negative prompt: {negative_prompt}")
-    print(f"Using seed: {current_seed}")
+    print(f"Using seed: {seed}") # Use the seed passed from the event handler
 
     pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
+    # Use the generator created with the potentially randomized seed
+    generator = torch.cuda.manual_seed(seed) if torch.cuda.is_available() else torch.manual_seed(seed)
+
     with torch.no_grad():
         initial_result = pipe(
             prompt=full_prompt,
             negative_prompt=negative_prompt,
             guidance_scale=guidance_scale,
-            num_images_per_prompt=1,
+            num_images_per_prompt=1, # Pipeline always generates 1 image here
             num_inference_steps=num_steps,
-            generator=generator,
+            generator=generator, # Pass the generator
             image=control_image,
         ).images[0]
 
@@ -337,83 +525,80 @@ def process_image_api(
         print(f"Error saving or uploading image: {e}")
 
     return initial_result
-
-# Define a Pydantic model for the request body parameters (optional, but good practice)
-# class ImageParameters(BaseModel):
-#     style_selection: str = "None"
-#     prompt: str = ""
-#     a_prompt: str = ""
-#     n_prompt: str = "EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)"
-#     image_resolution: int = 512
-#     preprocess_resolution: int = 512
-#     num_steps: int = 15
-#     guidance_scale: float = 5.5
-#     seed: int = -1
-
-# Define the API endpoint
-@app.post("/generate-image/")
-async def generate_image(
-    file: UploadFile = File(...), # Input image file
-    style_selection: str = Form("None"), # Parameters from form data
-    prompt: str = Form(""),
-    a_prompt: str = Form(""),
-    n_prompt: str = Form("EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)"),
-    image_resolution: int = Form(512),
-    preprocess_resolution: int = Form(512),
-    num_steps: int = Form(15),
-    guidance_scale: float = Form(5.5),
-    seed: int = Form(-1),
-):
-    """
-    API endpoint to generate an interior design image based on an input image and parameters.
-
-    Expects a POST request with form-data including:
-    - file: The input image file (UploadFile).
-    - style_selection: The design style name (string).
-    - prompt: Custom design prompt (string).
-    - a_prompt: Additional positive prompt (string).
-    - n_prompt: Negative prompt (string).
-    - image_resolution: Output image resolution (int).
-    - preprocess_resolution: Preprocessor resolution (int).
-    - num_steps: Number of inference steps (int).
-    - guidance_scale: Guidance scale (float).
-    - seed: Random seed (int, use -1 for random).
-
-    Returns:
-        The generated image as a JPEG file.
-    """
-    try:
-        # Read the uploaded image file
-        image_data = await file.read()
-        input_image = Image.open(io.BytesIO(image_data)).convert("RGB")
-
-        # Process the image using the core logic
-        generated_image = process_image_api(
-            image=input_image,
-            style_selection=style_selection,
-            prompt=prompt,
-            a_prompt=a_prompt,
-            n_prompt=n_prompt,
-            image_resolution=image_resolution,
-            preprocess_resolution=preprocess_resolution,
-            num_steps=num_steps,
-            guidance_scale=guidance_scale,
-            seed=seed,
-        )
-
-        # Return the generated image as a streaming response
-        buffer = io.BytesIO()
-        generated_image.save(buffer, format="JPEG")
-        buffer.seek(0)
-
-        return StreamingResponse(buffer, media_type="image/jpeg")
-
-    except Exception as e:
-        print(f"An error occurred during processing: {e}")
-        raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
-
-# Entry point to run the FastAPI application using Uvicorn
-if __name__ == "__main__":
-    # The host "0.0.0.0" makes the server accessible externally within the container
-    # The port is taken from the environment variable PORT, which Hugging Face Spaces sets
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    # --- End of process_image logic ---
+
+# Re-adding randomize_seed_fn
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+# Re-adding randomize_seed to config list
+config = [
+    image,
+    style_selection,
+    prompt,
+    a_prompt,
+    n_prompt,
+    num_images,
+    image_resolution,
+    preprocess_resolution,
+    num_steps,
+    guidance_scale,
+    seed,
+    randomize_seed, # Added randomize_seed to config
+]
+
+# Re-defining Gradio Event Handling Functions with randomize_seed input and seed randomization logic
+@gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config,
+       outputs=result, show_progress="minimal")
+def auto_process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images,
+                       image_resolution, preprocess_resolution, num_steps, guidance_scale,
+                       seed, randomize_seed, progress=gr.Progress(track_tqdm=True)): # Added randomize_seed
+    # Apply seed randomization
+    processed_seed = randomize_seed_fn(seed, randomize_seed)
+    print(f"Using processed seed: {processed_seed}") # Debug print
+
+    # Call the core processing function with the processed seed
+    return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images,
+                         image_resolution, preprocess_resolution, num_steps, guidance_scale,
+                         processed_seed) # Pass processed_seed
+
+@gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result],
+       show_progress="minimal")
+def submit(previous_result, image, style_selection, prompt, a_prompt, n_prompt, num_images,
+           image_resolution, preprocess_resolution, num_steps, guidance_scale, seed,
+           randomize_seed, progress=gr.Progress(track_tqdm=True)): # Added randomize_seed
+    # First, yield the previous result to update the input image immediately
+    yield previous_result, gr.update()
+
+    # Apply seed randomization
+    processed_seed = randomize_seed_fn(seed, randomize_seed)
+    print(f"Using processed seed: {processed_seed}") # Debug print
+
+    # Then, process the new input image
+    new_result = process_image(previous_result, style_selection, prompt, a_prompt,
+                               n_prompt, num_images, image_resolution,
+                               preprocess_resolution, num_steps, guidance_scale,
+                               processed_seed) # Pass processed_seed
+    # Finally, yield the new result
+    yield previous_result, new_result
+
+# Turn off buttons when processing - These functions remain the same
+@gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None,
+       outputs=[run_button, use_ai_button], show_progress="hidden")
+def turn_buttons_off():
+    return gr.update(visible=False), gr.update(visible=False)
+
+# Turn on buttons when processing is complete - These functions remain the same
+@gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button],
+       show_progress="hidden")
+def turn_buttons_on():
+    return gr.update(visible=True), gr.update(visible=True)
+
+
+# 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) # Use share=True for Spaces