Update app.py

app.py

@@ -21,11 +21,50 @@ pipe = None
 bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary", "death", "ghost"]'))
 default_negative = os.getenv("default_negative","")
 
+# Remove the GPU-specific import and decorator
+# import spaces
+# @spaces.GPU(enable_queue=True)
+
+# Update the check_text function to be more informative
 def check_text(prompt, negative=""):
-    for i in bad_words:
-        if i in prompt:
-            return True
-    return False
+    restricted_words = []
+    for word in bad_words:
+        if word in prompt.lower() or word in negative.lower():
+            restricted_words.append(word)
+    return restricted_words
+
+
+    restricted_words = check_text(prompt, negative_prompt)
+    if restricted_words:
+        return [], seed, f"Prompt contains restricted words: {', '.join(restricted_words)}"
+
+    # Ensure prompt is 2-3 words long
+    prompt = " ".join(prompt.split()[:3])
+
+    # Apply style
+    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.manual_seed(seed)
+
+    width, height = size_map.get(size, (512, 512))
+
+    if not use_negative_prompt:
+        negative_prompt = ""
+
+    options = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": 20,  # Reduced from 25
+        "generator": generator,
+        "num_images_per_prompt": 2,  # Reduced from 6
+        "output_type": "pil",
+    }
+    
+
+
 
 # Kid-friendly styles
 style_list = [
@@ -125,6 +164,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 @spaces.GPU(enable_queue=True)
+# Update the generate function
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -135,7 +175,6 @@ def generate(
     guidance_scale: float = 3,
     randomize_seed: bool = False,
     background: str = "transparent",
-    device_type: str = "cpu",
     progress=gr.Progress(track_tqdm=True),
 ):
     global device, pipe
@@ -177,7 +216,8 @@ def generate(
     images = pipe(**options).images
     image_paths = [save_image(img, background) for img in images]
     
-    return image_paths, seed
+    return image_paths, seed, None  # Added None for potential error message
+    
 
 examples = [
     "cute bunny",
@@ -189,15 +229,9 @@ css = '''
 .gradio-container{max-width: 700px !important}
 h1{text-align:center}
 '''
-
-# Define the Gradio UI for the sticker generator
+# Update the Gradio interface
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
-    )
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -209,6 +243,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             )
             run_button = gr.Button("Run")
         result = gr.Gallery(label="Generated Stickers", columns=2, preview=True)
+        error_output = gr.Textbox(label="Error", visible=False)
     with gr.Accordion("Advanced options", open=False):
         use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
         negative_prompt = gr.Text(
@@ -258,7 +293,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed],
+        outputs=[result, seed, error_output],
         fn=generate,
         cache_examples=CACHE_EXAMPLES,
     )
@@ -280,9 +315,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             randomize_seed,
             background_selection,
-            device_selection,
         ],
-        outputs=[result, seed],
+        outputs=[result, seed, error_output],
         api_name="run",
     )