Update app.py

app.py

@@ -11,10 +11,19 @@ from PIL import Image
 import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
+import logging
 
-# Force CPU usage
-device = torch.device("cpu")
-torch.cuda.is_available = lambda: False
+# 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"]'))
@@ -68,18 +77,32 @@ DESCRIPTION = """## Children's Sticker Generator
 
 Generate fun and playful stickers for children using AI.
 
-⚠️ Running on CPU. This may be slower.
 """
+DESCRIPTION += "🚀 Running on GPU for faster generation." if device.type == "cuda" else "⚠️ Running on CPU. This may be slower."
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = False
+CACHE_EXAMPLES = True
 
 # Initialize the DiffusionPipeline
-pipe = DiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",  # Using a smaller model for CPU
-    torch_dtype=torch.float32,
-    use_safetensors=True,
-).to(device)
+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
 
 # Convert mm to pixels for a specific DPI (300) and ensure divisible by 8
 def mm_to_pixels(mm, dpi=300):
@@ -127,7 +150,7 @@ def generate(
     style: str = DEFAULT_STYLE_NAME,
     seed: int = 0,
     size: str = "75mm",
-    guidance_scale: float = 3,
+    guidance_scale: float = 7.5,
     randomize_seed: bool = False,
     background: str = "transparent",
     progress=gr.Progress(track_tqdm=True),
@@ -142,9 +165,9 @@ def generate(
     # Apply style
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    width, height = size_map.get(size, (512, 512))
+    width, height = size_map.get(size, (1024, 1024))
 
     if not use_negative_prompt:
         negative_prompt = ""
@@ -155,9 +178,9 @@ def generate(
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": 20,
+        "num_inference_steps": 30 if device.type == "cuda" else 20,
         "generator": generator,
-        "num_images_per_prompt": 1,  # Reduced to 1 for CPU
+        "num_images_per_prompt": 4 if device.type == "cuda" else 1,
         "output_type": "pil",
     }
     
@@ -196,7 +219,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 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)
         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)
@@ -233,7 +256,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         )
         guidance_scale = gr.Slider(
             label="Guidance Scale",
-            minimum=0.1,
+            minimum=1.0,
             maximum=20.0,
             step=0.1,
             value=7.5,
@@ -270,4 +293,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    try:
+        demo.queue(max_size=20).launch()
+    except Exception as e:
+        logger.error(f"Error launching Gradio interface: {str(e)}")
+        raise