Update app.py

app.py

@@ -5,19 +5,23 @@ from PIL import Image
 import random
 from diffusers import DiffusionPipeline
 
-# Initialize DiffusionPipeline 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):
-    generated_img = pipeline(prompt)
-    return generated_img
+    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 = output.images[0]
+    generated_img_array = np.array(generated_img)
+    return generated_img_array
 
 def create_cereal_box(input_image):
-    cover_img = input_image  # This should already be a PIL Image
-    template_img = Image.open('CerealBoxMaker/template.jpeg')
-    
-    # Cereal box creation logic
+    cover_img = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    template_img = Image.open('/content/866b9b8f50b50879120be0b87dfd6050.jpg')
     scaling_factor = 1.5
     rect_height = int(template_img.height * 0.32)
     new_width = int(rect_height * 0.70)
@@ -34,15 +38,32 @@ 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 to a numpy array for Gradio output
     template_copy_array = np.array(template_copy)
     return template_copy_array
 
 def combined_function(prompt):
-    generated_img = text_to_image(prompt)
-    final_img = create_cereal_box(generated_img)
+    generated_img_array = text_to_image(prompt)
+    final_img = create_cereal_box(generated_img_array)
     return final_img
 
-# Create Gradio Interface
-gr.Interface(fn=combined_function, inputs="text", outputs="image").launch()
+with gr.Blocks() as app:
+    gr.HTML("<div style='text-align: center;'><h1>Cereal Box Maker 🥣</h1></div>")
+    gr.HTML("<div style='text-align: center;'><p>This application uses StableDiffusion XL to create any cereal box you could ever imagine!</p></div>")
+    gr.HTML("<div style='text-align: center;'><h3>Instructions:</h3><ol><li>Describe the cereal box you want to create and hit generate!</li><li>Print it out, cut the outside, fold the lines, and then tape!</li></ol></div>")
+    gr.HTML("<div style='text-align: center;'><p>A space by AP 🐧, follow me on <a href='https://twitter.com/angrypenguinPNG'>Twitter</a>! H/T to OstrisAI <a href='https://twitter.com/ostrisai'>Twitter</a> for their Cereal Box LoRA!</p></div>")
+    
+    with gr.Row():
+        textbox = gr.Textbox(label="Describe your cereal box: Ex: 'Avengers Cereal'")
+        btn_generate = gr.Button("Generate", label="Generate")
+    
+    with gr.Row():
+        output_img = gr.Image(label="Your Custom Cereal Box")
+
+    btn_generate.click(
+        combined_function,
+        inputs=[textbox],
+        outputs=[output_img]
+    )
+
+app.queue(concurrency_count=4, max_size=20, api_open=False)
+app.launch(debug=True)