File size: 6,172 Bytes
bb167dc a8f966d bb167dc 003ec88 bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc a8f966d bb167dc |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
from diffusers import StableDiffusionInpaintPipeline,StableDiffusionPipeline
from PIL import Image
import requests
import cv2
import torch
import matplotlib.pyplot as plt
import io
import requests
import os
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import gradio as gr
import gc
processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
IPmodel_path = "runwayml/stable-diffusion-inpainting"
IPpipe = StableDiffusionInpaintPipeline.from_pretrained(
IPmodel_path,
use_auth_token= st.secrets["AUTH_TOKEN"]
).to(device)
trans_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
trans_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
SDpipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token=st.secrets["AUTH_TOKEN"]).to(device)
def create_mask(image, prompt):
inputs = processor(text=[prompt], images=[image], padding="max_length", return_tensors="pt")
# predict
with torch.no_grad():
outputs = model(**inputs)
preds = outputs.logits
filename = f"mask.png"
plt.imsave(filename,torch.sigmoid(preds))
gray_image = cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)
(thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
# For debugging only:
# cv2.imwrite(filename,bw_image)
# fix color format
cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
mask = cv2.bitwise_not(bw_image)
cv2.imwrite(filename, mask)
return Image.open('mask.png')
def generate_image(image, product_name, target_name):
mask = create_mask(image, product_name)
image = image.resize((512, 512))
mask = mask.resize((512,512))
guidance_scale=8
#guidance_scale=16
num_samples = 4
prompt = target_name
generator = torch.Generator(device=device).manual_seed(22) # change the seed to get different results
im = IPpipe(
prompt=prompt,
image=image,
mask_image=mask,
guidance_scale=guidance_scale,
generator=generator,
).images
im.enable_attention_slicing()
return im
gc.collect()
torch.cuda.empty_cache()
image_label = 'Please upload the image (optional)'
extract_label = 'Specify what needs to be extracted from the above image'
prompt_label = 'Specify the description of image to be generated'
button_label = "Proceed"
output_label = "Results"
shot_services = ['close-up', 'extreme-closeup', 'POV','medium', 'long']
shot_label = 'Choose the shot type'
style_services = ['polaroid', 'monochrome', 'long exposure','color splash', 'Tilt shift']
style_label = 'Choose the style type'
lighting_services = ['soft', 'ambivalent', 'ring','sun', 'cinematic']
lighting_label = 'Choose the lighting type'
context_services = ['indoor', 'outdoor', 'at night','in the park', 'in the beach','studio']
context_label = 'Choose the context'
lens_services = ['wide angle', 'telephoto', '24 mm','EF 70mm', 'Bokeh']
lens_label = 'Choose the lens type'
device_services = ['iphone', 'CCTV', 'Nikon ZFX','Canon', 'Gopro']
device_label = 'Choose the device type'
def add_to_prompt(prompt_text,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio ):
if shot_radio != '':
prompt_text += ", "+shot_radio
if style_radio != '':
prompt_text += ", "+style_radio
if lighting_radio != '':
prompt_text += ", "+lighting_radio
if context_radio != '':
prompt_text += ", "+ context_radio
if lens_radio != '':
prompt_text += ", "+ lens_radio
if device_radio != '':
prompt_text += ", "+ device_radio
return prompt_text
def proceed_with_generation(input_file, extract_text, prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio):
if extract_text == "" or input_file == "":
prompt = add_to_prompt(prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
print(prompt)
output = SDpipe(prompt, height=512, width=512, num_images_per_prompt=4)
return output.images
elif extract_text != "" and input_file != "" and prompt_text !='':
prompt = add_to_prompt(prompt_text,shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio)
print(prompt)
print(extract_text)
output = generate_image(Image.fromarray(input_file), extract_text, prompt)
return output
else:
raise gr.Error("Please fill all details for guided image or atleast prompt for free image rendition !")
with gr.Blocks() as demo:
with gr.Row():
input_file = gr.Image(interactive = True, label=image_label, visible=False, shape=(512,512))
extract_text = gr.Textbox(label= extract_label, lines=1, interactive = True, visible = True)
prompt_text = gr.Textbox(label= prompt_label, lines=1, interactive = True, visible = True)
with gr.Accordion("Advanced Options", open=False):
shot_radio = gr.Radio(shot_services , label=shot_label, )
style_radio = gr.Radio(style_services , label=style_label)
lighting_radio = gr.Radio(lighting_services , label=lighting_label)
context_radio = gr.Radio(context_services , label=context_label)
lens_radio = gr.Radio(lens_services , label=lens_label)
device_radio = gr.Radio(device_services , label=device_label)
button = gr.Button(value = button_label , visible = True)
with gr.Row():
output_gallery = gr.Gallery(label = output_label, visible= True)
button.click( proceed_with_generation, [input_file, extract_text, prompt_text, shot_radio, style_radio, lighting_radio, context_radio, lens_radio, device_radio], [output_gallery])
demo.launch(debug=True) |