Update app.py

app.py

@@ -11,13 +11,13 @@ MODEL_NAME = "Heartsync/NSFW-Uncensored"
 CACHE_DIR = "./model_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Fixed device handling - removed invalid device_map
+# Load model with strict CPU settings
 pipe = StableDiffusionPipeline.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float32,  # Use FP32 for CPU
+    torch_dtype=torch.float32,
     cache_dir=CACHE_DIR,
     low_cpu_mem_usage=True,
-    safety_checker=None  # Optional: disable safety checker for CPU speed
+    safety_checker=None
 )
 
 DISCLAIMER = """
@@ -26,10 +26,19 @@ This model may generate explicit/sensitive content. By using this application, y
 """
 
 def generate_image(prompt, width, height, num_inference_steps=20, guidance_scale=7.0):
+    """Generate image with strict validation"""
     if not prompt.strip():
-        return None, "Please enter a prompt"
+        return None, "Error: Prompt cannot be empty"
     
     try:
+        # Force integer dimensions
+        width = int(width) if width else 512
+        height = int(height) if height else 512
+        
+        # Validate dimensions (must be divisible by 8)
+        width = width - (width % 8)
+        height = height - (height % 8)
+        
         with torch.inference_mode():
             result = pipe(
                 prompt=prompt,
@@ -38,10 +47,16 @@ def generate_image(prompt, width, height, num_inference_steps=20, guidance_scale
                 num_inference_steps=num_inference_steps,
                 guidance_scale=guidance_scale
             ).images[0]
-            return result.convert("RGB"), None
+            
+            if result is None:
+                return None, "Error: Failed to generate image"
+                
+            return result.convert("RGB"), f"Success ({width}x{height})"
+            
     except Exception as e:
-        return None, f"Error: {str(e)}"
+        return None, f"Generation Error: {str(e)}"
 
+# UI Configuration
 ASPECT_RATIOS = {
     "Square (512x512)": (512, 512),
     "Portrait (512x768)": (512, 768),
@@ -55,7 +70,11 @@ with gr.Blocks(theme="soft", css=".disclaimer {color: #FF4B2B; font-size: 0.9em;
     
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", placeholder="Describe your image...", lines=3)
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe your image...",
+                lines=3
+            )
             with gr.Row():
                 width = gr.Number(value=512, label="Width", interactive=True)
                 height = gr.Number(value=512, label="Height", interactive=True)
@@ -68,11 +87,11 @@ with gr.Blocks(theme="soft", css=".disclaimer {color: #FF4B2B; font-size: 0.9em;
         
         with gr.Column():
             output_image = gr.Image(label="Result", interactive=False)
-            error_msg = gr.Textbox(label="Status", visible=True)
+            error_msg = gr.Textbox(label="Status", value="Ready", interactive=False)
     
     def update_dimensions(choice):
         if choice == "Custom":
-            return [gr.Number(interactive=True), gr.Number(interactive=True)]
+            return [gr.Number(value=512, interactive=True), gr.Number(value=512, interactive=True)]
         else:
             w, h = ASPECT_RATIOS[choice]
             return [gr.Number(value=w, interactive=True), gr.Number(value=h, interactive=True)]