Update app.py

app.py

@@ -1,12 +1,11 @@
 # Configuration
-prod = False # This variable is no longer used for launching, but kept for potential future use
-port = 8080 # This variable is no longer used for launching, but kept for potential future use
-show_options = False # This variable is no longer used for UI visibility
+# These variables are now mostly for reference, FastAPI/Uvicorn handle port
+prod = False
+port = int(os.environ.get("PORT", 8080)) # Use PORT environment variable provided by Spaces, default to 8080
 
 import os
 import random
 import time
-# Removed gradio import as UI is being removed
 import numpy as np
 import spaces
 import imageio
@@ -19,20 +18,27 @@ 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
+
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 # os.environ['HF_HOME'] = '/data/.huggingface'
 print("CUDA version:", torch.version.cuda)
 print("loading everything")
-compiled = False # This variable is no longer explicitly set to True after compilation print
+compiled = False
 api = HfApi()
 
+# Initialize FastAPI app
+app = FastAPI()
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
@@ -46,8 +52,8 @@ class Preprocessor:
             return
         elif name == "NormalBae":
             print("Loading NormalBae")
-            # Ensure model is moved to cuda if available
-            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda" if torch.cuda.is_available() else "cpu")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to(device)
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
             self.name = name
@@ -56,24 +62,22 @@ class Preprocessor:
         return
 
     def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
-        # Ensure model is on the correct device before calling
         device = "cuda" if torch.cuda.is_available() else "cpu"
         if self.model.device.type != device:
              self.model.to(device)
         return self.model(image, **kwargs)
 
-# Load models and preprocessor directly without gr.NO_RELOAD check
-# This block will execute when the script is imported or run
+# Load models and preprocessor when the script starts
 # Controlnet Normal
-model_id = "lllyasviel/control_v11p_sd15_normalbae"
+model_id = "lllylyasviel/control_v11p_sd15_normalbae" # Corrected model ID based on common usage
 print("initializing controlnet")
-# Ensure models are loaded onto the correct device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 controlnet = ControlNetModel.from_pretrained(
     model_id,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Use float32 if CUDA is not available
-    attn_implementation="flash_attention_2" if torch.cuda.is_available() else None, # Flash attention only for CUDA
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    attn_implementation="flash_attention_2" if torch.cuda.is_available() else None,
 ).to(device)
+
 # Scheduler
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
     "ashllay/stable-diffusion-v1-5-archive",
@@ -85,21 +89,19 @@ scheduler = DPMSolverMultistepScheduler.from_pretrained(
     prediction_type="epsilon",
     thresholding=False,
     denoise_final=True,
-    # device_map="cuda", # device_map can sometimes cause issues, better to move after loading
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Use float32 if CUDA is not available
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
 )
-# Removed this line as Schedulers don't have a .to() method
-# scheduler.to(device)
+# Schedulers do not need to be moved to device
 
 # Stable Diffusion Pipeline URL
 base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
 print('loading pipe')
 pipe = StableDiffusionControlNetPipeline.from_single_file(
     base_model_url,
-    safety_checker=None,
+    safety_checker=None, # Keep None for now, but consider enabling for public API
     controlnet=controlnet,
     scheduler=scheduler,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Use float32 if CUDA is not available
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
 ).to(device)
 
 print("loading preprocessor")
@@ -108,68 +110,37 @@ preprocessor.load("NormalBae") # Preprocessor is loaded here
 
 # Load textual inversions
 try:
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20",
-        weight_name="HDA_unconventional maid.pt",
-        token="HDA_unconventional_maid",
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress"
-    )
-    pipe.load_textual_inversion(
-        "broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari"
-    )
+    # List of textual inversions to load
+    textual_inversions = {
+        "EasyNegativeV2": "EasyNegativeV2.safetensors",
+        "badhandv4": "badhandv4.pt",
+        "fcNeg-neg": "fcNeg-neg.pt",
+        "HDA_Ahegao": "HDA_Ahegao.pt",
+        "HDA_Bondage": "HDA_Bondage.pt",
+        "HDA_pet_play": "HDA_pet_play.pt",
+        "HDA_unconventional_maid": "HDA_unconventional maid.pt",
+        "HDA_NakedHoodie": "HDA_NakedHoodie.pt",
+        "HDA_NunDress": "HDA_NunDress.pt",
+        "HDA_Shibari": "HDA_Shibari.pt",
+    }
+    for token, weight_name in textual_inversions.items():
+        try:
+            pipe.load_textual_inversion(
+                "broyang/hentaidigitalart_v20", weight_name=weight_name, token=token,
+            )
+            print(f"Loaded textual inversion: {token}")
+        except Exception as e:
+            print(f"Warning: Could not load textual inversion {weight_name}: {e}")
+
 except Exception as e:
-    print(f"Error loading textual inversions: {e}")
-    # Handle cases where loading textual inversions might fail, e.g., file not found
+    print(f"Error during textual inversions loading process: {e}")
+
 
 print("---------------Loaded controlnet pipeline---------------")
 if torch.cuda.is_available():
     torch.cuda.empty_cache()
     gc.collect()
     print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
-# Removed "Model Compiled!" print as compilation is not explicitly handled here
-
-
-# Removed generate_furniture_mask as inpainting part was commented out
-# def generate_furniture_mask(image, furniture_type):
-#     image_np = np.array(image)
-#     height, width = image_np.shape[:2]
-#     mask = np.zeros((height, width), dtype=np.uint8)
-#     if furniture_type == "sofa":
-#         cv2.rectangle(mask, (width // 4, int(height * 0.6)), (width * 3 // 4, height), 255, -1)
-#     elif furniture_type == "table":
-#         cv2.rectangle(mask, (width // 3, height // 3), (width * 2 // 3, height * 2 // 3), 255, -1)
-#     elif furniture_type == "chair":
-#         cv2.circle(mask, (width * 3 // 5, height * 2 // 3), height // 6, 255, -1)
-#     return Image.fromarray(mask)
-
-# Removed randomize_seed_fn as the logic is directly in process_image
-# def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-#     return seed
 
 
 def get_additional_prompt():
@@ -184,51 +155,29 @@ def get_additional_prompt():
 
 def get_prompt(prompt, additional_prompt):
     interior = "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"
-    # default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed" # Not used
-    # default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed" # Not used
-    randomize = get_additional_prompt()
-    # nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed" # Not used
-    # bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW" # Not used
-    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
-    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
-    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
-    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao" # Not used
-    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
-    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
-    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
-    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
-    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
-    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
-    abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
-    # shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari" # Not used
-    shibari2 = "octane render, highly detailed, volumetric, HDA_Shibari"
-
-    if prompt == "":
-        # This block seems to generate prompts for 'girls' which might not be relevant for interior design API
-        # Consider if this random girl prompt generation is needed for the interior design API
-        girls = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress,
-                 naked_hoodie, abg, shibari2, ahegao2]
-        # prompts_nsfw = [abg, shibari2, ahegao2] # Not used
-        # prompt = f"{random.choice(girls)}" # This line would overwrite the input prompt
-        prompt = f"boho chic" # This line also overwrites the input prompt
-        # The logic here seems inconsistent with using an input 'prompt'.
-        # Assuming the intention is to use the input 'prompt' for custom designs,
-        # and apply a style or default interior context.
-        # Let's revise this to prioritize the input prompt.
-        if additional_prompt:
-             # Combine input prompt with additional prompt
-             return f"{prompt}, {additional_prompt}"
-        else:
-             # If no additional prompt, just use the input prompt with interior context
-             return f"Photo from Pinterest of {prompt} {interior}"
+
+    # Revised logic to prioritize the input prompt and combine with interior context and additional prompt
+    prompt_parts = []
+    if prompt:
+        prompt_parts.append(f"Photo from Pinterest of {prompt}")
     else:
-        # If a prompt is provided, use it with the interior context
-        # The original logic here was redundant with the 'if prompt == ""' block
-        # Let's simplify based on whether a prompt is provided
-        if additional_prompt:
-            return f"Photo from Pinterest of {prompt} {interior}, {additional_prompt}"
-        else:
-            return f"Photo from Pinterest of {prompt} {interior}"
+        # 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
+
+    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))
 
 
 style_list = [
@@ -285,31 +234,22 @@ STYLE_NAMES = list(styles.keys())
 
 
 def apply_style(style_name):
-    # Ensure style_name exists in styles dictionary
-    return styles.get(style_name, "") # Return empty string if style not found
-
+    return styles.get(style_name, "")
 
-# Removed css variable as it was for Gradio UI
-# css = """..."""
 
-# Removed gr.Blocks context manager and everything inside it
-
-
-# Modified process_image to be a standalone function callable by an API endpoint
-# Removed @spaces.GPU and @torch.inference_mode decorators if the API framework handles this
-# Added type hints for clarity
+# 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)",
-    num_images: int = 1, # Kept for potential future use, but pipeline currently generates 1
     image_resolution: int = 512,
     preprocess_resolution: int = 512,
     num_steps: int = 15,
     guidance_scale: float = 5.5,
-    seed: int = -1, # Use -1 to indicate random seed if not provided
+    seed: int = -1,
 ):
     """
     Processes an input image to generate a new image based on style and prompts.
@@ -320,7 +260,6 @@ def process_image_api(
         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.
@@ -330,18 +269,14 @@ def process_image_api(
     Returns:
         A PIL Image of the generated result.
     """
-    # Use provided seed or generate a random one
     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)
 
-    # Ensure preprocessor is loaded
     if preprocessor.name != "NormalBae":
          preprocessor.load("NormalBae")
 
-    # Ensure preprocessor model is on the correct device
     preprocessor.model.to("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Generate control image
     control_image = preprocessor(
         image=image,
         image_resolution=image_resolution,
@@ -350,50 +285,38 @@ def process_image_api(
 
     # Construct the full prompt
     if style_selection and style_selection != "None":
-        # Apply selected style and combine with custom prompt and additional prompt
         style_prompt = apply_style(style_selection)
-        # Combine prompts, ensuring no empty strings lead to awkward commas
         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:
-        # Use custom prompt and additional prompt with default interior context
         full_prompt = get_prompt(prompt, a_prompt)
 
+
     negative_prompt = str(n_prompt)
     print(f"Using prompt: {full_prompt}")
     print(f"Using negative prompt: {negative_prompt}")
     print(f"Using seed: {current_seed}")
 
-
-    # Generate the image using the pipeline
-    # Ensure the pipeline is on the correct device
     pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-    with torch.no_grad(): # Use no_grad for inference to save memory and speed
+    with torch.no_grad():
         initial_result = pipe(
             prompt=full_prompt,
             negative_prompt=negative_prompt,
             guidance_scale=guidance_scale,
-            num_images_per_prompt=1, # Pipeline always generates 1 image here
+            num_images_per_prompt=1,
             num_inference_steps=num_steps,
             generator=generator,
             image=control_image,
         ).images[0]
 
-    # Save and upload results (optional, depending on API requirements)
-    # This part might be handled by the API caller or a separate service
-    # Keeping it for now as it was in the original script
+    # Save and upload results (optional)
     try:
         timestamp = int(time.time())
-        # Saving input image is generally not needed for API response, but keeping for consistency
-        # img_path = f"{timestamp}_input.jpg"
         results_path = f"{timestamp}_output.jpg"
-        # imageio.imsave(img_path, image) # Removed saving input image
         imageio.imsave(results_path, initial_result)
 
-        # Uploading files might not be desired for a general API,
-        # consider making this optional or removing if the API just returns the image
-        if API_KEY: # Only attempt upload if API_KEY is available
+        if API_KEY:
             print(f"Uploading result image to broyang/interior-ai-outputs/{results_path}")
             try:
                 api.upload_file(
@@ -402,17 +325,8 @@ def process_image_api(
                     repo_id="broyang/interior-ai-outputs",
                     repo_type="dataset",
                     token=API_KEY,
-                    run_as_future=True, # Asynchronous upload
+                    run_as_future=True,
                 )
-                # Removed input image upload
-                # api.upload_file(
-                #     path_or_fileobj=img_path,
-                #     path_in_repo=img_path,
-                #     repo_id="broyang/interior-ai-outputs",
-                #     repo_type="dataset",
-                #     token=API_KEY,
-                #     run_as_future=True,
-                # )
             except Exception as e:
                 print(f"Error uploading file to Hugging Face Hub: {e}")
         else:
@@ -421,32 +335,85 @@ def process_image_api(
     except Exception as e:
         print(f"Error saving or uploading image: {e}")
 
-
     return initial_result
 
-# The script now defines the process_image_api function.
-# To use this as an API, you would typically import this script
-# into a web framework like FastAPI and define an endpoint that
-# calls process_image_api with the appropriate parameters from the request.
-
-# Example of how you might call the function (this part is for demonstration,
-# you would remove it when integrating into a web framework):
-# if __name__ == "__main__":
-#     # Create a dummy input image (e.g., a black square)
-#     dummy_image = Image.new('RGB', (512, 512), color = 'red')
-#     print("Generating a sample image...")
-#     # Call the processing function with sample parameters
-#     generated_image = process_image_api(
-#         image=dummy_image,
-#         style_selection="Boho",
-#         prompt="cozy living room",
-#         a_prompt="warm lighting",
-#         num_steps=20,
-#         guidance_scale=7.0,
-#         seed=42
-#     )
-#     # You can now save or display the generated_image
-#     generated_image.save("sample_output.jpg")
-#     print("Sample image generated and saved as sample_output.jpg")
-
-# Removed the demo.queue().launch() calls
+# 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)
+