Update app.py

app.py

@@ -8,24 +8,17 @@ import re
 import gradio as gr
 import numpy as np
 from PIL import Image
-import spaces
 import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
 
-# Initialize device to None
-device = None
-pipe = None
+# Check if GPU is available; fallback to CPU if needed
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Setup rules for bad words (ensure the prompts are kid-friendly)
 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=""):
     restricted_words = []
     for word in bad_words:
@@ -33,39 +26,6 @@ def check_text(prompt, negative=""):
             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 = [
     {
@@ -108,21 +68,18 @@ DESCRIPTION = """## Children's Sticker Generator
 Generate fun and playful stickers for children using AI.
 """
 
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>⚠️Running on CPU, This may be slower.</p>"
+
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
 
-def initialize_pipeline(device_type):
-    global device, pipe
-    device = torch.device(device_type)
-    pipe = DiffusionPipeline.from_pretrained(
-        "SG161222/RealVisXL_V3.0_Turbo",
-        torch_dtype=torch.float32 if device_type == "cpu" else torch.float16,
-        use_safetensors=True,
-    ).to(device)
-
-# Initialize with CPU by default
-initialize_pipeline("cpu")
-
+# Initialize the DiffusionPipeline
+pipe = DiffusionPipeline.from_pretrained(
+    "SG161222/RealVisXL_V3.0_Turbo",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    use_safetensors=True,
+).to(device)
 
 # Convert mm to pixels for a specific DPI (300) and ensure divisible by 8
 def mm_to_pixels(mm, dpi=300):
@@ -163,8 +120,6 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU(enable_queue=True)
-# Update the generate function
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -177,28 +132,22 @@ def generate(
     background: str = "transparent",
     progress=gr.Progress(track_tqdm=True),
 ):
-    global device, pipe
-    
-    # Switch device if necessary
-    if device.type != device_type:
-        initialize_pipeline(device_type)
+    restricted_words = check_text(prompt, negative_prompt)
+    if restricted_words:
+        return [], seed, f"Prompt contains restricted words: {', '.join(restricted_words)}"
 
-    if check_text(prompt, negative_prompt):
-        raise ValueError("Prompt contains restricted words.")
-    
     # Ensure prompt is 2-3 words long
-    prompt = " ".join(re.findall(r'\w+', prompt)[:3])
+    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.Generator(device=device).manual_seed(seed)
 
-    # Ensure we have only white or transparent background options
-    width, height = size_map.get(size, (1024, 1024))
+    width, height = size_map.get(size, (512, 512))
 
     if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
+        negative_prompt = ""
 
     options = {
         "prompt": prompt,
@@ -206,9 +155,9 @@ def generate(
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": 25,
+        "num_inference_steps": 20,
         "generator": generator,
-        "num_images_per_prompt": 6,  # Max 6 images
+        "num_images_per_prompt": 2,
         "output_type": "pil",
     }
     
@@ -216,8 +165,7 @@ def generate(
     images = pipe(**options).images
     image_paths = [save_image(img, background) for img in images]
     
-    return image_paths, seed, None  # Added None for potential error message
-    
+    return image_paths, seed, None
 
 examples = [
     "cute bunny",
@@ -229,9 +177,15 @@ css = '''
 .gradio-container{max-width: 700px !important}
 h1{text-align:center}
 '''
-# Update the Gradio interface
+
+# Define the Gradio UI for the sticker generator
 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(
@@ -284,11 +238,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             step=0.1,
             value=15.7,
         )
-        device_selection = gr.Radio(
-            choices=["cpu", "cuda"],
-            value="cpu",
-            label="Device",
-        )
 
     gr.Examples(
         examples=examples,