Update app.py

app.py

@@ -8,33 +8,20 @@ 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
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Check for GPU availability and fall back to CPU if necessary
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    logger.info("GPU detected. Using CUDA.")
-else:
-    device = torch.device("cpu")
-    logger.warning("No GPU detected. Falling back to CPU. This will be slower.")
 
 # 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","")
 
 def check_text(prompt, negative=""):
-    restricted_words = []
-    for word in bad_words:
-        if word in prompt.lower() or word in negative.lower():
-            restricted_words.append(word)
-    return restricted_words
+    for i in bad_words:
+        if i in prompt:
+            return True
+    return False
 
 # Kid-friendly styles
 style_list = [
@@ -76,33 +63,23 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
 DESCRIPTION = """## Children's Sticker Generator
 
 Generate fun and playful stickers for children using AI.
-
 """
-DESCRIPTION += "🚀 Running on GPU for faster generation." if device.type == "cuda" else "⚠️ Running on CPU. This may be slower."
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = True
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Initialize the DiffusionPipeline
-try:
-    if device.type == "cuda":
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16",
-        ).to(device)
-        pipe.enable_xformers_memory_efficient_attention()
-    else:
-        pipe = DiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
-            torch_dtype=torch.float32,
-            use_safetensors=True,
-        ).to(device)
-    logger.info("DiffusionPipeline initialized successfully")
-except Exception as e:
-    logger.error(f"Error initializing DiffusionPipeline: {str(e)}")
-    raise
+pipe = DiffusionPipeline.from_pretrained(
+    "SG161222/RealVisXL_V3.0_Turbo",  # or any model of your choice
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+).to(device)
 
 # Convert mm to pixels for a specific DPI (300) and ensure divisible by 8
 def mm_to_pixels(mm, dpi=300):
@@ -143,6 +120,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+@spaces.GPU(enable_queue=True)
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -150,27 +128,27 @@ def generate(
     style: str = DEFAULT_STYLE_NAME,
     seed: int = 0,
     size: str = "75mm",
-    guidance_scale: float = 7.5,
+    guidance_scale: float = 3,
     randomize_seed: bool = False,
     background: str = "transparent",
     progress=gr.Progress(track_tqdm=True),
 ):
-    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(prompt.split()[:3])
+    prompt = " ".join(re.findall(r'\w+', prompt)[: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)
+    generator = torch.Generator().manual_seed(seed)
 
+    # Ensure we have only white or transparent background options
     width, height = size_map.get(size, (1024, 1024))
 
     if not use_negative_prompt:
-        negative_prompt = ""
+        negative_prompt = ""  # type: ignore
 
     options = {
         "prompt": prompt,
@@ -178,9 +156,9 @@ def generate(
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": 30 if device.type == "cuda" else 20,
+        "num_inference_steps": 25,
         "generator": generator,
-        "num_images_per_prompt": 4 if device.type == "cuda" else 1,
+        "num_images_per_prompt": 6,  # Max 6 images
         "output_type": "pil",
     }
     
@@ -188,7 +166,7 @@ def generate(
     images = pipe(**options).images
     image_paths = [save_image(img, background) for img in images]
     
-    return image_paths, seed, None
+    return image_paths, seed
 
 examples = [
     "cute bunny",
@@ -220,7 +198,6 @@ 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(
@@ -256,16 +233,16 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         )
         guidance_scale = gr.Slider(
             label="Guidance Scale",
-            minimum=1.0,
+            minimum=0.1,
             maximum=20.0,
             step=0.1,
-            value=7.5,
+            value=6,
         )
 
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed, error_output],
+        outputs=[result, seed],
         fn=generate,
         cache_examples=CACHE_EXAMPLES,
     )
@@ -288,13 +265,9 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             randomize_seed,
             background_selection,
         ],
-        outputs=[result, seed, error_output],
+        outputs=[result, seed],
         api_name="run",
     )
 
 if __name__ == "__main__":
-    try:
-        demo.queue(max_size=20).launch()
-    except Exception as e:
-        logger.error(f"Error launching Gradio interface: {str(e)}")
-        raise
+    demo.queue(max_size=20).launch()