Update app.py

app.py

@@ -2,29 +2,37 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
+import random
 from diffusers import DiffusionPipeline
 
 # Initialize the DiffusionPipeline model with LoRA weights
 pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
 pipeline.load_lora_weights("ostris/super-cereal-sdxl-lora")
+pipeline.to("cuda:0")
+
+MAX_SEED = np.iinfo(np.int32).max
 
 def text_to_image(prompt):
-    # Generate image using the DiffusionPipeline
-    output = pipeline(prompt)
+    seed = random.randint(0, MAX_SEED)
+    negative_prompt = "ugly, blurry, nsfw, gore, blood"
+    
+    output = pipeline(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=1024,
+        height=1024,
+        guidance_scale=7.0,
+        num_inference_steps=25,
+        generator=torch.Generator().manual_seed(seed),
+    )
+    
     generated_img_tensor = output.images[0]
-
-    # Convert torch tensor to numpy array
     generated_img_array = generated_img_tensor.cpu().numpy().transpose((1, 2, 0))
     return generated_img_array
 
 def create_cereal_box(input_image):
-    # Convert the input numpy array to PIL Image
     cover_img = Image.fromarray((input_image.astype(np.uint8)))
-
-    # Load the template image
-    template_img = Image.open('CerealBoxMaker/template.jpeg') 
-
-    # Simplified cereal box creation logic
+    template_img = Image.open('CerealBoxMaker/template.jpeg')
     scaling_factor = 1.5
     rect_height = int(template_img.height * 0.32)
     new_width = int(rect_height * 0.70)
@@ -41,10 +49,7 @@ def create_cereal_box(input_image):
     template_copy = template_img.copy()
     template_copy.paste(cover_resized_scaled, left_position)
     template_copy.paste(cover_resized_scaled, right_position)
-
-    # Convert the PIL Image back to a numpy array
     template_copy_array = np.array(template_copy)
-    
     return template_copy_array
 
 def combined_function(prompt):
@@ -53,4 +58,4 @@ def combined_function(prompt):
     return final_img
 
 # Create a Gradio Interface
-gr.Interface(fn=combined_function, inputs="text", outputs="image").launch()
+gr.Interface(fn=combined_function, inputs="text", outputs="image").launch(debug=True)