Raushan-123 commited on
Commit
297632a
·
verified ·
1 Parent(s): 562e1af

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import DiffusionModel, DiffusionImageProcessor, AutoTokenizer
4
+ from threading import Thread
5
+
6
+ print("Starting to load the model to memory")
7
+ # Load the diffusion model and image processor
8
+ model = DiffusionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
9
+ processor = DiffusionImageProcessor.from_model(model)
10
+ tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
11
+
12
+ print("Successfully loaded the model to memory")
13
+
14
+ def generate_image(text):
15
+ # Generate an image from the given text prompt
16
+ inputs = tokenizer(text, return_tensors="pt")
17
+ # Run generation on GPU if available
18
+ inputs.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
19
+
20
+ # Run diffusion model for image generation
21
+ with torch.no_grad():
22
+ result = processor.generate(**inputs)
23
+
24
+ # Return the generated image
25
+ return result[0]
26
+
27
+ # Define a function to handle user input and generate images
28
+ def image_generator(text):
29
+ generated_image = generate_image(text)
30
+ return generated_image
31
+
32
+ # Create a Gradio interface for the image generation
33
+ interface = gr.Interface(
34
+ fn=image_generator,
35
+ inputs="text",
36
+ outputs="image",
37
+ title="Image Generation from Text",
38
+ description="Enter a text prompt to generate an image.",
39
+ examples=["a cat sitting on a couch"]
40
+ )
41
+
42
+ # Launch the interface
43
+ interface.launch(share=True)