Create app.py

app.py

@@ -0,0 +1,139 @@
+import gradio as gr
+import spaces
+from diffusers import AutoPipelineForText2Image
+import torch
+import time
+# import logging
+from threading import Timer
+from nsfw_detector import NSFWDetector
+
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
+
+# Global variables
+pipe = None
+last_use_time = None
+unload_timer = None
+TIMEOUT_SECONDS = 120  # 2 minutes
+BATCH_SIZE = 4
+
+def chunk_generations(num_images):
+    """Split number of images into batches of BATCH_SIZE"""
+    return [min(BATCH_SIZE, num_images - i) for i in range(0, num_images, BATCH_SIZE)]
+
+@spaces.GPU(duration=20)
+def generate_image(
+    prompt,
+    num_inference_steps=1,
+    num_images=1,
+    height=512,
+    width=512,
+):    
+    global pipe
+    start_time = time.time()
+    
+    # Load model if needed
+    if pipe is None:
+        yield None, "Loading model..."
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            "stabilityai/sdxl-turbo", 
+            torch_dtype=torch.float16, 
+            variant="fp16"
+        ).to("cuda")
+        yield None, "Model loaded, starting generation..."
+    
+    reset_timer()
+    
+    # Process in batches if more than BATCH_SIZE images
+    if num_images > BATCH_SIZE:
+        yield None, f"Generating {num_images} images in batches..."
+        all_images = []
+        batches = chunk_generations(num_images)
+        
+        for i, batch_size in enumerate(batches):
+            yield None, f"Generating batch {i+1}/{len(batches)} ({batch_size} images)..."
+            batch_images = pipe(
+                prompt=prompt,
+                num_inference_steps=num_inference_steps,
+                height=height,
+                width=width,
+                guidance_scale=0.0,
+                num_images_per_prompt=batch_size
+            ).images
+            all_images.extend(batch_images)
+        images = all_images
+    else:
+        yield None, f"Generating {num_images} image(s) with {num_inference_steps} steps..."
+        images = pipe(
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            height=height,
+            width=width,
+            guidance_scale=0.0,
+            num_images_per_prompt=num_images
+        ).images
+    
+    total_time = time.time() - start_time
+    avg_time = total_time / num_images
+    status_msg = f"Generated {num_images} image(s) in {total_time:.2f} seconds (avg {avg_time:.2f}s per image)"
+    # logger.info(status_msg)
+    
+    # Check for NSFW content
+    detector = NSFWDetector()
+    is_nsfw, category, confidence = detector.check_image(images[0])
+        
+    if category == "SAFE":
+        yield images, status_msg
+    else:
+        return
+
+def unload_model():
+    global pipe, last_use_time
+    current_time = time.time()
+    if last_use_time and (current_time - last_use_time) >= TIMEOUT_SECONDS:
+        # logger.info("Unloading model due to inactivity...")
+        pipe = None
+        torch.cuda.empty_cache()
+        return "Model unloaded due to inactivity"
+
+def reset_timer():
+    global unload_timer, last_use_time
+    if unload_timer:
+        unload_timer.cancel()
+    last_use_time = time.time()
+    unload_timer = Timer(TIMEOUT_SECONDS, unload_model)
+    unload_timer.start()
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    with gr.Column():
+        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
+        with gr.Row():
+            steps = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Number of inference steps")
+            num_images = gr.Slider(minimum=1, maximum=64, value=1, step=1, label="Number of images to generate")
+        with gr.Row():
+            height = gr.Slider(minimum=512, maximum=1024, value=512, step=64, label="Height")
+            width = gr.Slider(minimum=512, maximum=1024, value=512, step=64, label="Width")
+        
+        generate_btn = gr.Button("Generate")
+        gallery = gr.Gallery()
+        status = gr.Textbox(
+            label="Status", 
+            value="Model not loaded - will load on first generation", 
+            interactive=False
+        )
+        
+        generate_btn.click(
+            fn=generate_image,
+            inputs=[prompt, steps, num_images, height, width],
+            outputs=[gallery, status]
+        )
+
+    gr.Markdown("""
+        This model works best with 512x512 resolution and 1-4 inference steps. 
+        Values above 4 steps may not improve quality significantly.
+        The model will automatically unload after 2 minutes of inactivity.
+    """)
+
+if __name__ == "__main__":
+    demo.launch()