File size: 1,916 Bytes
12863e1
2cf6be0
 
fb01197
05d3d42
987f112
 
 
 
2cf6be0
4095388
fb01197
92ecdcd
fb01197
92ecdcd
4095388
 
6b0d828
2cf6be0
 
 
12863e1
4095388
 
987f112
0bbe8f6
987f112
 
 
 
 
2cf6be0
fa0ee64
 
 
 
fb01197
92ecdcd
f0f8ecd
262e1d3
987f112
92ecdcd
ad9ba71
92ecdcd
2cf6be0
fa0ee64
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import spaces
import gradio as gr
import torch
from diffusers import DiffusionPipeline, AutoencoderKL
import rembg
from io import BytesIO
import PIL.Image as Image
import cv2
import numpy

model_id = "dataautogpt3/OpenDalleV1.1"
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = DiffusionPipeline.from_pretrained(model_id,
                                         vae=vae,
                                         torch_dtype=torch.float16,
                                        use_safetensors=True,
                                        variant="fp16")

pipe.to("cuda")

# Function to generate an image from text using diffusion
@spaces.GPU
def generate_image(prompt, neg_prompt):
    prompt += "no background, side view, minimalist shot"
    image_bytes = pipe(prompt, negative_prompt=neg_prompt)
    print(image_bytes)
    pil_image = Image.open(BytesIO(image_bytes))
    opencvImage = cv2.cvtColor(numpy.array(pil_image), cv2.COLOR_RGB2BGR)

    cv2_image = rembg.remove(opencvImage)
    return pil_image, cv2_image

_TITLE = "Shoe Generator"
with gr.Blocks(_TITLE) as ShoeGen:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Enter a discription of a shoe")
            neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
            button_gen = gr.Button("Generate Image")
        with gr.Column():
            image = gr.Image(label="Generated Image")   
            image2 = gr.Image(label="Generated Image without background", show_download_button=True) 
    
    button_gen.click(generate_image, inputs=[prompt], outputs=[image, image2])

ShoeGen.launch()