Kvikontent commited on
Commit
8406177
Β·
1 Parent(s): 6226130

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
4
+ from PIL import Image
5
+ import requests
6
+ from io import BytesIO
7
+
8
+ # Set up the Stable Diffusion pipeline for text to image
9
+ text_to_image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
10
+ text_to_image_pipe = text_to_image_pipe.to("cuda")
11
+
12
+ # Set up the Stable Diffusion pipeline for image to image
13
+ image_to_image_pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
14
+ device = "cuda"
15
+ image_to_image_pipe = image_to_image_pipe.to(device)
16
+
17
+ def text_to_image(prompt):
18
+ image = text_to_image_pipe(prompt).images[0]
19
+ image = Image.fromarray(image)
20
+ return image
21
+
22
+ def image_to_image(file, prompt):
23
+ response = requests.get(file)
24
+ init_image = Image.open(BytesIO(response.content)).convert("RGB")
25
+ init_image = init_image.resize((768, 512))
26
+ images = image_to_image_pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
27
+ return images[0]
28
+
29
+ # Create the Gradio interface
30
+ image_to_image_interface = gr.inputs.Image(label="Input Image")
31
+ text_prompt_interface = gr.inputs.Textbox(label="Text Prompt")
32
+
33
+ io_choice = gr.dropdown(["Text to Image", "Image to Image"], label="Choose Input Type")
34
+
35
+ outputs = []
36
+
37
+ if io_choice.lower() == "text to image":
38
+ outputs.append(gr.outputs.Image(label="Generated Image"))
39
+ interface = gr.Interface(fn=text_to_image, inputs=text_prompt_interface, outputs=outputs)
40
+ else:
41
+ outputs.append(gr.outputs.Image(label="Generated Image"))
42
+ interface = gr.Interface(fn=image_to_image, inputs=[image_to_image_interface, text_prompt_interface], outputs=outputs)
43
+
44
+ # Launch the interface
45
+ interface.launch()