Update app.py

app.py

@@ -4,7 +4,7 @@ import os
 import random
 import uuid
 import json
-
+import re
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -13,25 +13,18 @@ import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
 
-# Check for the Model Base..//
-
-bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary"]'))
-bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
+# 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=""):
     for i in bad_words:
         if i in prompt:
             return True
-    for i in bad_words_negative:
-        if i in negative:
-            return True
     return False
 
-# Updated to child-friendly styles
-
+# Kid-friendly styles
 style_list = [
-
     {
         "name": "Cartoon",
         "prompt": "colorful cartoon {prompt}. vibrant, playful, friendly, suitable for children, highly detailed, bright colors",
@@ -42,19 +35,16 @@ style_list = [
         "prompt": "children's illustration {prompt}. cute, colorful, fun, simple shapes, smooth lines, highly detailed, joyful",
         "negative_prompt": "scary, dark, violent, deformed, ugly",
     },   
-
     {
         "name": "Sticker",
         "prompt": "children's sticker of {prompt}. bright colors, playful, high resolution, cartoonish",
         "negative_prompt": "scary, dark, violent, ugly, low resolution",
     },
-
     {
         "name": "Fantasy",
         "prompt": "fantasy world for children with {prompt}. magical, vibrant, friendly, beautiful, colorful",
         "negative_prompt": "dark, scary, violent, ugly, realistic",
     },
-
     {
         "name": "(No style)",
         "prompt": "{prompt}",
@@ -68,8 +58,6 @@ DEFAULT_STYLE_NAME = "Sticker"
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    if not negative:
-        negative = ""
     return p.replace("{prompt}", positive), n + negative
 
 DESCRIPTION = """## Children's Sticker Generator
@@ -82,43 +70,39 @@ if not torch.cuda.is_available():
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-NUM_IMAGES_PER_PROMPT = 1
-
-if torch.cuda.is_available():
-    pipe = DiffusionPipeline.from_pretrained(
-        "SG161222/RealVisXL_V3.0_Turbo",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        add_watermarker=False,
-        variant="fp16"
-    )
-    pipe2 = DiffusionPipeline.from_pretrained(
-        "SG161222/RealVisXL_V2.02_Turbo",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        add_watermarker=False,
-        variant="fp16"
-    )
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-        pipe2.enable_model_cpu_offload()
-    else:
-        pipe.to(device)    
-        pipe2.to(device)    
-        print("Loaded on Device!")
-    
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-
-def save_image(img):
+NUM_IMAGES_PER_PROMPT = 6  # Set maximum images per prompt
+
+# Convert mm to pixels for a specific DPI (300)
+def mm_to_pixels(mm, dpi=300):
+    return int((mm / 25.4) * dpi)
+
+# Default sizes for 75mm and 35mm
+size_map = {
+    "75mm": (mm_to_pixels(75), mm_to_pixels(75)),  # 75mm in pixels at 300dpi
+    "35mm": (mm_to_pixels(35), mm_to_pixels(35)),  # 35mm in pixels at 300dpi
+}
+
+# Function to post-process images (transparent or white background)
+def save_image(img, background="transparent"):
+    img = img.convert("RGBA")
+    data = img.getdata()
+    new_data = []
+
+    if background == "transparent":
+        for item in data:
+            # Replace white with transparent
+            if item[0] == 255 and item[1] == 255 and item[2] == 255:
+                new_data.append((255, 255, 255, 0))  # Transparent
+            else:
+                new_data.append(item)
+    elif background == "white":
+        for item in data:
+            new_data.append(item)  # Keep as white
+
+    img.putdata(new_data)
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
@@ -135,23 +119,28 @@ def generate(
     use_negative_prompt: bool = False,
     style: str = DEFAULT_STYLE_NAME,
     seed: int = 0,
-    width: int = 512,
-    height: int = 512,
+    size: str = "75mm",
     guidance_scale: float = 3,
     randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
+    background: str = "transparent",
     progress=gr.Progress(track_tqdm=True),
 ):
     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])
+
+    # Apply style
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_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 = ""  # type: ignore
-    negative_prompt += default_negative    
 
     options = {
         "prompt": prompt,
@@ -162,19 +151,19 @@ def generate(
         "num_inference_steps": 25,
         "generator": generator,
         "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,
-        "use_resolution_binning": use_resolution_binning,
         "output_type": "pil",
     }
     
-    images = pipe(**options).images + pipe2(**options).images
-
-    image_paths = [save_image(img) for img in images]
+    # Generate images with the pipeline
+    images = pipe(**options).images
+    image_paths = [save_image(img, background) for img in images]
+    
     return image_paths, seed
 
 examples = [
-    "A cute cartoon bunny holding a carrot in a colorful garden",
-    "A playful dragon flying through the clouds, bright and friendly",
-    "A magical unicorn standing on a rainbow with sparkles",
+    "cute bunny",
+    "happy cat",
+    "funny dog",
 ]
 
 css = '''
@@ -182,6 +171,7 @@ css = '''
 h1{text-align:center}
 '''
 
+# Define the Gradio UI for the sticker generator
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(
@@ -195,36 +185,20 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 label="Enter your prompt",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter a fun, child-friendly idea (e.g., cute bunny with a rainbow)",
+                placeholder="Enter 2-3 word prompt (e.g., cute bunny)",
                 container=False,
             )
             run_button = gr.Button("Run")
-        result = gr.Gallery(label="Generated Stickers", columns=1, preview=True)
+        result = gr.Gallery(label="Generated Stickers", columns=2, preview=True)
     with gr.Accordion("Advanced options", open=False):
         use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
         negative_prompt = gr.Text(
             label="Negative prompt",
             max_lines=1,
             placeholder="Enter a negative prompt",
-            value="(scary, violent, deformed, ugly, dark)",
+            value="(scary, violent, dark, ugly)",
             visible=True,
         )
-        with gr.Row():
-            num_inference_steps = gr.Slider(
-                label="Steps",
-                minimum=10,
-                maximum=60,
-                step=1,
-                value=25,
-            )
-        with gr.Row():
-            num_images_per_prompt = gr.Slider(
-                label="Images",
-                minimum=1,
-                maximum=5,
-                step=1,
-                value=2,
-            )
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -234,38 +208,24 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             visible=True
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row(visible=True):
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=1024,
-                step=8,
-                value=512,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=1024,
-                step=8,
-                value=512,
-            )
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance Scale",
-                minimum=0.1,
-                maximum=20.0,
-                step=0.1,
-                value=7,
-            )
-    with gr.Row(visible=True):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=STYLE_NAMES,
-            value=DEFAULT_STYLE_NAME,
-            label="Sticker Style",
+        size_selection = gr.Radio(
+            choices=["75mm", "35mm"],
+            value="75mm",
+            label="Sticker Size",
+        )
+        background_selection = gr.Radio(
+            choices=["transparent", "white"],
+            value="transparent",
+            label="Background Color",
+        )
+        guidance_scale = gr.Slider(
+            label="Guidance Scale",
+            minimum=0.1,
+            maximum=20.0,
+            step=0.1,
+            value=6,
         )
+
     gr.Examples(
         examples=examples,
         inputs=prompt,
@@ -274,13 +234,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         cache_examples=CACHE_EXAMPLES,
     )
 
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        api_name=False,
-    )
-
     gr.on(
         triggers=[
             prompt.submit,
@@ -294,10 +247,10 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             use_negative_prompt,
             style_selection,
             seed,
-            width,
-            height,
+            size_selection,
             guidance_scale,
             randomize_seed,
+            background_selection,
         ],
         outputs=[result, seed],
         api_name="run",