123LETSPLAY commited on
Commit
a3ea6c4
1 Parent(s): eb55fa7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from diffusers import DiffusionPipeline
3
+ import gradio as gr
4
+
5
+ # Load the diffusion pipeline using the CogVideoX model
6
+ pipe = DiffusionPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.float16)
7
+
8
+ def generate_image(prompt):
9
+ # Generate an image based on the prompt using the CogVideoX model
10
+ image = pipe(prompt).images[0]
11
+ return image
12
+
13
+ # Set up the Gradio interface
14
+ with gr.Blocks() as demo:
15
+ gr.Markdown("# Image Generation using CogVideoX Diffusers")
16
+
17
+ # Input for user's text prompt
18
+ prompt_input = gr.Textbox(label="Enter Prompt", value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
19
+
20
+ # Output for displaying the generated image
21
+ output_image = gr.Image(label="Generated Image")
22
+
23
+ # Button to trigger image generation
24
+ generate_button = gr.Button("Generate Image")
25
+
26
+ # Link button click with image generation function
27
+ generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
28
+
29
+ # Launch the Gradio app
30
+ demo.launch()
31
+