Update app.py

app.py

@@ -2,70 +2,98 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 import os
+from PIL import Image
+import warnings
 
-# Set environment variables for CPU optimization
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-os.environ["HF_HOME"] = "./cache"
+# Suppress warnings for cleaner output
+warnings.filterwarnings("ignore")
 
-# Load model with memory-efficient settings
+# Model loading configuration
+MODEL_NAME = "Heartsync/NSFW-Uncensored"
+CACHE_DIR = "./model_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# Load model with memory optimization
 pipe = StableDiffusionPipeline.from_pretrained(
-    "Heartsync/NSFW-Uncensored",
-    torch_dtype=torch.float32,  # Required for CPU
+    MODEL_NAME,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    cache_dir=CACHE_DIR,
+    revision="fp16" if torch.cuda.is_available() else None,
     low_cpu_mem_usage=True,
-    use_safetensors=True,
-    cache_dir="./model_cache"
-).to("cpu")
+    device_map="cpu"
+)
 
-# Enable memory-saving techniques
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
+# Safety disclaimer
+DISCLAIMER = """
+⚠️ **Content Warning**  
+This model may generate explicit/sensitive content. By using this application, you agree to comply with [HuggingFace's Acceptable Use Policy](https://huggingface.co/policies/acceptable-use ).  
+"""
 
-def generate_image(prompt, negative_prompt="", num_inference_steps=30, guidance_scale=7.5, num_images=1):
-    """Generate images with NSFW model, ensuring safe memory usage"""
+def generate_image(prompt, width, height, num_inference_steps=30, guidance_scale=7.5):
+    """Generate image with safety checks and memory management"""
+    if not prompt.strip():
+        return None, "Please enter a prompt"
+    
     try:
-        images = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images,
-            generator=torch.Generator(device="cpu").manual_seed(42)
-        ).images
-        
-        return images
+        with torch.inference_mode():
+            result = pipe(
+                prompt=prompt,
+                width=width,
+                height=height,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale
+            ).images[0]
+            return result.convert("RGB"), None
     except Exception as e:
-        return f"Error: {str(e)}. This model may require shorter prompts or lower resolution."
+        return None, f"Error: {str(e)}"
 
-# Define UI
-with gr.Blocks(theme="soft", css=".gr-box {border-color: #ff4b5c}") as demo:
-    gr.Markdown("""
-    # 🔞 NSFW-Uncensored Image Generator
-    ⚠️ This model generates explicit content. Use responsibly and in compliance with local laws.
-    
-    Model: [Heartsync/NSFW-Uncensored](https://huggingface.co/Heartsync/NSFW-Uncensored )
-    """)
+# UI Configuration
+ASPECT_RATIOS = {
+    "Square (512x512)": (512, 512),
+    "Portrait (512x768)": (512, 768),
+    "Landscape (768x512)": (768, 512),
+    "Phone (384x832)": (384, 832),
+    "Custom": None
+}
+
+with gr.Blocks(theme="soft", css=".disclaimer {color: #FF4B2B; font-size: 0.9em;}") as demo:
+    gr.Markdown("# 🔞 NSFW-Uncensored Text-to-Image Generator\n" + DISCLAIMER)
     
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", placeholder="Describe your image...", lines=4)
-            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What to avoid...", lines=3)
-            
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe your image...",
+                lines=3
+            )
             with gr.Row():
-                num_images = gr.Slider(1, 4, value=1, step=1, label="Images to Generate")
-                num_inference_steps = gr.Slider(20, 50, value=30, step=1, label="Generation Steps")
-            
-            guidance_scale = gr.Slider(1, 20, value=7.5, step=0.5, label="Guidance Scale")
-            
+                width = gr.Number(value=512, label="Width", interactive=True)
+                height = gr.Number(value=512, label="Height", interactive=True)
+            aspect_ratio = gr.Dropdown(
+                choices=list(ASPECT_RATIOS.keys()),
+                value="Square (512x512)",
+                label="Aspect Ratio"
+            )
+            generate_btn = gr.Button("Generate Image", variant="primary")
+        
         with gr.Column():
-            output = gr.Gallery(label="Generated Images", columns=2, height=512)
+            output_image = gr.Image(label="Result", interactive=False)
+            error_msg = gr.Textbox(label="Status", visible=True)
+    
+    def update_dimensions(choice):
+        if choice == "Custom":
+            return [gr.Number(interactive=True), gr.Number(interactive=True)]
+        else:
+            w, h = ASPECT_RATIOS[choice]
+            return [gr.Number(value=w, interactive=True), gr.Number(value=h, interactive=True)]
     
-    generate_btn = gr.Button("🎨 Generate NSFW Image", variant="primary")
+    aspect_ratio.change(fn=update_dimensions, inputs=[aspect_ratio], outputs=[width, height])
     
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, num_images],
-        outputs=output
+        inputs=[prompt, width, height],
+        outputs=[output_image, error_msg]
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.launch()