Update app.py

app.py

@@ -1,14 +1,12 @@
-prod = False
-port = 8080
-show_options = False
-if prod:
-    port = 8081
-    # show_options = False
+# 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
 
 import os
 import random
 import time
-import gradio as gr
+# Removed gradio import as UI is being removed
 import numpy as np
 import spaces
 import imageio
@@ -21,20 +19,21 @@ from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
-    # StableDiffusionInpaintPipeline,
-    # AutoencoderKL,
+    # 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
 
 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
+compiled = False # This variable is no longer explicitly set to True after compilation print
 api = HfApi()
 
+
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
 
@@ -47,128 +46,153 @@ class Preprocessor:
             return
         elif name == "NormalBae":
             print("Loading NormalBae")
-            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
-            torch.cuda.empty_cache()
+            # 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")
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
             self.name = name
         else:
             raise ValueError
         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)
 
-if gr.NO_RELOAD:
-    # Controlnet Normal
-    model_id = "lllyasviel/control_v11p_sd15_normalbae"
-    print("initializing controlnet")
-    controlnet = ControlNetModel.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        attn_implementation="flash_attention_2",
-    ).to("cuda")
-
-    # Scheduler
-    scheduler = DPMSolverMultistepScheduler.from_pretrained(
-        # "runwayml/stable-diffusion-v1-5",
-        # "stable-diffusion-v1-5/stable-diffusion-v1-5",
-        "ashllay/stable-diffusion-v1-5-archive",
-        solver_order=2,
-        subfolder="scheduler",
-        use_karras_sigmas=True,
-        final_sigmas_type="sigma_min",
-        algorithm_type="sde-dpmsolver++",
-        prediction_type="epsilon",
-        thresholding=False,
-        denoise_final=True,
-        device_map="cuda",
-        torch_dtype=torch.float16,
+# Load models and preprocessor directly without gr.NO_RELOAD check
+# This block will execute when the script is imported or run
+# Controlnet Normal
+model_id = "lllyasviel/control_v11p_sd15_normalbae"
+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
+).to(device)
+# Scheduler
+scheduler = DPMSolverMultistepScheduler.from_pretrained(
+    "ashllay/stable-diffusion-v1-5-archive",
+    solver_order=2,
+    subfolder="scheduler",
+    use_karras_sigmas=True,
+    final_sigmas_type="sigma_min",
+    algorithm_type="sde-dpmsolver++",
+    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
+)
+# Move scheduler to device after loading
+scheduler.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,
+    controlnet=controlnet,
+    scheduler=scheduler,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Use float32 if CUDA is not available
+).to(device)
+
+print("loading preprocessor")
+preprocessor = Preprocessor()
+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"
+    )
+except Exception as e:
+    print(f"Error loading textual inversions: {e}")
+    # Handle cases where loading textual inversions might fail, e.g., file not found
 
-    # Stable Diffusion Pipeline URL
-    # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
-    base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
-    # vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
-
-    # print('loading vae')
-    # vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
-    # vae.to(memory_format=torch.channels_last) 
-
-    print('loading pipe')
-    pipe = StableDiffusionControlNetPipeline.from_single_file(
-        base_model_url,
-        safety_checker=None,
-        controlnet=controlnet,
-        scheduler=scheduler,
-        # vae=vae,
-        torch_dtype=torch.float16,
-    ).to("cuda")
-    
-    # print('loading inpainting pipe')
-    # inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    #     "runwayml/stable-diffusion-inpainting",
-    #     torch_dtype=torch.float16,
-    # ).to("cuda")
-
-    print("loading preprocessor")
-    preprocessor = Preprocessor()
-    preprocessor.load("NormalBae")
-    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")
-    pipe.to("cuda")
-
-    print("---------------Loaded controlnet pipeline---------------") 
+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")
-    print("Model Compiled!")
-    
-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)
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+# 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():
     prompt = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     top = ["tank top", "blouse", "button up shirt", "sweater", "corset top"]
-    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt", "leggings", "high-waisted shorts"]
-    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker", "necklace", "headband", "headphones"]
+    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt",
+              "leggings", "high-waisted shorts"]
+    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker",
+                 "necklace", "headband", "headphones"]
     return f"{prompt}, {random.choice(top)}, {random.choice(bottom)}, {random.choice(accessory)}, score_9"
-    # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]
+
 
 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"
-    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    # 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"
-    # bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
+    # 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"
+    # 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"
@@ -176,330 +200,253 @@ def get_prompt(prompt, additional_prompt):
     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"
+    # 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 == "":
-        girls = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
-        prompts_nsfw = [abg, shibari2, ahegao2]
-        prompt = f"{random.choice(girls)}"
-        prompt = f"boho chic"
-        # print(f"-------------{preset}-------------")
+        # 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}"
     else:
-        prompt = f"Photo from Pinterest of {prompt} {interior}"
-        # prompt = default2
-    return f"{prompt} f{additional_prompt}"
+        # 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}"
 
-style_list = [
-    {
-        "name": "None",
-        "prompt": ""
-    },
-    {
-        "name": "Minimalistic",
-        "prompt": "Minimalist interior design,clean lines,neutral colors,uncluttered space,functional furniture,lots of natural light"
-    },
-    {
-        "name": "Boho",
-        "prompt": "Bohemian chic interior,eclectic mix of patterns and textures,vintage furniture,plants,woven textiles,warm earthy colors"
-    },
-    {
-        "name": "Farmhouse",
-        "prompt": "Modern farmhouse interior,rustic wood elements,shiplap walls,neutral color palette,industrial accents,cozy textiles"
-    },
-    {
-        "name": "Saudi Prince",
-        "prompt": "Opulent gold interior,luxurious ornate furniture,crystal chandeliers,rich fabrics,marble floors,intricate Arabic patterns"
-    },
-    {
-        "name": "Neoclassical",
-        "prompt": "Neoclassical interior design,elegant columns,ornate moldings,symmetrical layout,refined furniture,muted color palette"
-    },
-    {
-        "name": "Eclectic",
-        "prompt": "Eclectic interior design,mix of styles and eras,bold color combinations,diverse furniture pieces,unique art objects"
-    },
-    {
-        "name": "Parisian",
-        "prompt": "Parisian apartment interior,all-white color scheme,ornate moldings,herringbone wood floors,elegant furniture,large windows"
-    },
-    {
-        "name": "Hollywood",
-        "prompt": "Hollywood Regency interior,glamorous and luxurious,bold colors,mirrored surfaces,velvet upholstery,gold accents"
-    },
-    {
-        "name": "Scandinavian",
-        "prompt": "Scandinavian interior design,light wood tones,white walls,minimalist furniture,cozy textiles,hygge atmosphere"
-    },
-    {
-        "name": "Beach",
-        "prompt": "Coastal beach house interior,light blue and white color scheme,weathered wood,nautical accents,sheer curtains,ocean view"
-    },
-    {
-        "name": "Japanese",
-        "prompt": "Traditional Japanese interior,tatami mats,shoji screens,low furniture,zen garden view,minimalist decor,natural materials"
-    },
-    { 
-        "name": "Midcentury Modern",
-        "prompt": "Mid-century modern interior,1950s-60s style furniture,organic shapes,warm wood tones,bold accent colors,large windows"
-    },
-    {
-        "name": "Retro Futurism",
-        "prompt": "Neon (atompunk world) retro cyberpunk background",
-    },
-    {
-        "name": "Texan",
-        "prompt": "Western cowboy interior,rustic wood beams,leather furniture,cowhide rugs,antler chandeliers,southwestern patterns"
-    },
-    {
-        "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"
-    }
-] 
 
+style_list = [
+    {"name": "None",
+     "prompt": ""
+     },
+    {"name": "Minimalistic",
+     "prompt": "Minimalist interior design,clean lines,neutral colors,uncluttered space,functional furniture,lots of natural light"
+     },
+    {"name": "Boho",
+     "prompt": "Bohemian chic interior,eclectic mix of patterns and textures,vintage furniture,plants,woven textiles,warm earthy colors"
+     },
+    {"name": "Farmhouse",
+     "prompt": "Modern farmhouse interior,rustic wood elements,shiplap walls,neutral color palette,industrial accents,cozy textiles"
+     },
+    {"name": "Saudi Prince",
+     "prompt": "Opulent gold interior,luxurious ornate furniture,crystal chandeliers,rich fabrics,marble floors,intricate Arabic patterns"
+     },
+    {"name": "Neoclassical",
+     "prompt": "Neoclassical interior design,elegant columns,ornate moldings,symmetrical layout,refined furniture,muted color palette"
+     },
+    {"name": "Eclectic",
+     "prompt": "Eclectic interior design,mix of styles and eras,bold color combinations,diverse furniture pieces,unique art objects"
+     },
+    {"name": "Parisian",
+     "prompt": "Parisian apartment interior,all-white color scheme,ornate moldings,herringbone wood floors,elegant furniture,large windows"
+     },
+    {"name": "Hollywood",
+     "prompt": "Hollywood Regency interior,glamorous and luxurious,bold colors,mirrored surfaces,velvet upholstery,gold accents"
+     },
+    {"name": "Scandinavian",
+     "prompt": "Scandinavian interior design,light wood tones,white walls,minimalist furniture,cozy textiles,hygge atmosphere"
+     },
+    {"name": "Beach",
+     "prompt": "Coastal beach house interior,light blue and white color scheme,weathered wood,nautical accents,sheer curtains,ocean view"
+     },
+    {"name": "Japanese",
+     "prompt": "Traditional Japanese interior,tatami mats,shoji screens,low furniture,zen garden view,minimalist decor,natural materials"
+     },
+    {"name": "Midcentury Modern",
+     "prompt": "Mid-century modern interior,1950s-60s style furniture,organic shapes,warm wood tones,bold accent colors,large windows"
+     },
+    {"name": "Retro Futurism",
+     "prompt": "Neon (atompunk world) retro cyberpunk background",
+     },
+    {"name": "Texan",
+     "prompt": "Western cowboy interior,rustic wood beams,leather furniture,cowhide rugs,antler chandeliers,southwestern patterns"
+     },
+    {"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"
+     }]
 styles = {k["name"]: (k["prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
 
+
 def apply_style(style_name):
-    if style_name in styles:
-        p = styles.get(style_name, "none")
-    return p
-
-    
-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;
-}
-"""
-with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
-    #############################################################################
-    with gr.Row():
-        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
-            )  # 20/4.5 or 12 without lora, 4 with lora
-            guidance_scale = gr.Slider(
-                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
-            )  # 5 without lora, 2 with lora
-            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)
-    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)
-    
-    # image processing
-    @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)):
-        return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-    
-    # AI image processing
-    @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
-
-    # Turn off buttons when processing
-    @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
-    @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)
-
-@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,
+    # Ensure style_name exists in styles dictionary
+    return styles.get(style_name, "") # Return empty string if style not found
+
+
+# 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
+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 = random.randint(0, MAX_SEED)
-    generator = torch.cuda.manual_seed(seed)
-    
-    preprocessor.load("NormalBae")
+    """
+    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. Use -1 for random seed.
+
+    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,
         detect_resolution=preprocess_resolution,
     )
-    
-    if style_selection is not None and style_selection != "None":
-        prompt = f"Photo from Pinterest of {apply_style(style_selection)} {prompt},{a_prompt}"
+
+    # 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:
-        prompt = str(get_prompt(prompt, a_prompt))
+        # Use custom prompt and additional prompt with default interior context
+        full_prompt = get_prompt(prompt, a_prompt)
+
     negative_prompt = str(n_prompt)
-    print(prompt)
-    
-    # Generate the initial room image
-    initial_result = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=1,
-        num_inference_steps=num_steps,
-        generator=generator,
-        image=control_image,
-    ).images[0]
-    
-    # # Randomly choose whether to add furniture and which type
-    # furniture_types = ["sofa", "table", "chair", "dresser", "bookshelf", "desk", "coffee table"]
-    # furniture_type = random.choice(furniture_types)
-    
-    
-    # furniture_mask = generate_furniture_mask(initial_result, furniture_type)
-    # furniture_prompt = f"{prompt}, with a {furniture_type} in the style of {style_selection}"
-    # print(furniture_prompt)
-    
-    # # Use the inpainting pipeline to add furniture
-    # final_result = inpaint_pipe(
-    #     prompt=furniture_prompt,
-    #     image=initial_result,
-    #     mask_image=furniture_mask,
-    #     negative_prompt=negative_prompt,
-    #     num_inference_steps=num_steps,
-    #     guidance_scale=guidance_scale,
-    #     generator=generator,
-    # ).images[0]
-    
-    # Save and upload results
-    timestamp = int(time.time())
-    img_path = f"{timestamp}.jpg"
-    results_path = f"{timestamp}_out.jpg"
-    imageio.imsave(img_path, image)
-    imageio.imsave(results_path, initial_result)
-    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,
-    )
-    api.upload_file(
-        path_or_fileobj=results_path,
-        path_in_repo=results_path,
-        repo_id="broyang/interior-ai-outputs",
-        repo_type="dataset",
-        token=API_KEY,
-        run_as_future=True,
-    )
+    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
+        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_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
+    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
+            print(f"Uploading result image to broyang/interior-ai-outputs/{results_path}")
+            try:
+                api.upload_file(
+                    path_or_fileobj=results_path,
+                    path_in_repo=results_path,
+                    repo_id="broyang/interior-ai-outputs",
+                    repo_type="dataset",
+                    token=API_KEY,
+                    run_as_future=True, # Asynchronous upload
+                )
+                # 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:
+            print("Hugging Face API Key not found, skipping file upload.")
+
+    except Exception as e:
+        print(f"Error saving or uploading image: {e}")
+
+
     return initial_result
 
-if prod:
-    demo.queue(max_size=20).launch(server_name="localhost", server_port=port)
-else:
-    demo.queue(api_open=False).launch(show_api=False)
+# 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