yukee1992 commited on
Commit
db0fb8b
·
verified ·
1 Parent(s): f6b3207

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +57 -0
app.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Import necessary libraries
2
+ import torch
3
+ from diffusers import StableDiffusionXLPipeline
4
+ import gradio as gr
5
+ import time
6
+
7
+ # Check if we have a GPU (CUDA) or need to use the CPU
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+
10
+ # Load the Stable Diffusion XL pipeline
11
+ # We use torch_dtype=torch.float16 for faster generation and less memory usage
12
+ print("Loading Stable Diffusion XL pipeline... This may take a few minutes.")
13
+ pipe = StableDiffusionXLPipeline.from_pretrained(
14
+ "stabilityai/stable-diffusion-xl-base-1.0",
15
+ torch_dtype=torch.float16,
16
+ use_safetensors=True
17
+ )
18
+ # Move the pipeline to the chosen device (GPU or CPU)
19
+ pipe = pipe.to(device)
20
+ print("Model loaded successfully!")
21
+
22
+ # Define the image generation function
23
+ def generate_image(prompt):
24
+ """
25
+ This function takes a text prompt and returns a generated image.
26
+ """
27
+ # Add a consistent style to all prompts to get a children's book look
28
+ enhanced_prompt = f"children's book illustration, watercolor style, cute, whimsical, {prompt}"
29
+
30
+ # Generate the image with some default parameters
31
+ # guidance_scale controls how closely the image follows the prompt
32
+ image = pipe(
33
+ prompt=enhanced_prompt,
34
+ guidance_scale=9.5,
35
+ num_inference_steps=25 # More steps can mean higher quality, but is slower
36
+ ).images[0] # We get the first (and only) image from the result
37
+
38
+ return image
39
+
40
+ # Create the Gradio Interface
41
+ # We are creating a simple one-input (prompt) one-output (image) interface
42
+ demo = gr.Interface(
43
+ fn=generate_image, # The function to call
44
+ inputs=gr.Textbox( # The input is a text box for the prompt
45
+ label="Enter your scene description",
46
+ lines=2,
47
+ placeholder="A brave little mouse exploring a giant forest..."
48
+ ),
49
+ outputs=gr.Image(label="Generated Illustration", type="pil"), # The output is an image
50
+ title="Children's Book Illustrator 🤖🎨",
51
+ description="Generate beautiful illustrations for your children's story. Enter a description of a scene."
52
+ )
53
+
54
+ # This is the key part for making the API work correctly with n8n.
55
+ # We launch the Gradio app with a custom `api_name`.
56
+ # Setting `api_name="generate"` creates an API endpoint at `/api/generate/`
57
+ demo.launch(debug=True, server_name="0.0.0.0", api_name="generate")