Update generate function

app.py

@@ -6,7 +6,7 @@ import base64
 import torch
 from diffusers import StableDiffusionPipeline
 
-from transformers import pipeline
+#from transformers import pipeline
 import gradio as gr
 
 # Set Hugging Face API (needed for gated models)
@@ -16,7 +16,7 @@ hf_api_key = os.environ.get('HF_API_KEY')
 model_id = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionPipeline.from_pretrained(
     model_id,
-    torch_dtype=torch.float16,  # Use float16 for better performance on GPU
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 on GPU, float32 on CPU
     use_auth_token=hf_api_key  # Required for gated model
 )
 
@@ -41,18 +41,21 @@ pipe = pipe.to(device)
 #    return result_image
 
 def generate(prompt, negative_prompt, steps, guidance, width, height):
+    # Ensure width and height are multiples of 8 (required by Stable Diffusion)
+    width = int(width) - (int(width) % 8)
+    height = int(height) - (int(height) % 8)
+    
     # Generate image with Stable Diffusion
     output = pipe(
         prompt,
-        negative_prompt=negative_prompt,
+        negative_prompt=negative_prompt or None,  # Handle empty negative prompt
         num_inference_steps=int(steps),
         guidance_scale=float(guidance),
-        width=int(width),
-        height=int(height)
+        width=width,
+        height=height
     )
     return output.images[0]  # Return the first generated image (PIL format)
 
-
 # Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Generation with Stable Diffusion")