Spaces:
Runtime error
Runtime error
Commit
·
19aec01
1
Parent(s):
70138aa
Upload app.py
Browse files
app.py
ADDED
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
"""Untitled3.ipynb
|
3 |
+
|
4 |
+
Automatically generated by Colaboratory.
|
5 |
+
|
6 |
+
Original file is located at
|
7 |
+
https://colab.research.google.com/drive/1BltKPv_n-glCuuIIYSBA6GHK-tmwbl20
|
8 |
+
"""
|
9 |
+
|
10 |
+
!pip install fastai
|
11 |
+
!pip install gradio --upgrade
|
12 |
+
!pip install lida
|
13 |
+
!pip install diffusers --upgrade
|
14 |
+
!pip install tensorflow-probability --upgrade
|
15 |
+
!pip install invisible_watermark transformers accelerate safetensors
|
16 |
+
!pip install torch --upgrade
|
17 |
+
!pip install transformers --upgrade
|
18 |
+
!pip install datasets
|
19 |
+
#!pip install typing_extensions==4.5.0
|
20 |
+
#!pip install fastapi==0.103.0
|
21 |
+
|
22 |
+
import torch
|
23 |
+
from PIL import Image
|
24 |
+
from torchvision import transforms
|
25 |
+
import gradio as gr
|
26 |
+
import json
|
27 |
+
import urllib, urllib.request
|
28 |
+
from diffusers import DiffusionPipeline
|
29 |
+
from transformers import pipeline
|
30 |
+
|
31 |
+
# First Page
|
32 |
+
def demo_tab(image):
|
33 |
+
# The demo tab simply returns the same image as input
|
34 |
+
return image
|
35 |
+
|
36 |
+
# Secibd Page
|
37 |
+
def generate_image(Prompt,Negative_prompt,Steps):
|
38 |
+
# load both base & refiner
|
39 |
+
base = DiffusionPipeline.from_pretrained(
|
40 |
+
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
|
41 |
+
)
|
42 |
+
base.to("cuda")
|
43 |
+
refiner = DiffusionPipeline.from_pretrained(
|
44 |
+
"stabilityai/stable-diffusion-xl-refiner-1.0",
|
45 |
+
text_encoder_2=base.text_encoder_2,
|
46 |
+
vae=base.vae,
|
47 |
+
torch_dtype=torch.float16,
|
48 |
+
use_safetensors=True,
|
49 |
+
variant="fp16",
|
50 |
+
)
|
51 |
+
refiner.to("cuda")
|
52 |
+
|
53 |
+
# Define how many steps and what % of steps to be run on each experts (80/20) here
|
54 |
+
high_noise_frac = 0.8
|
55 |
+
prompt = Prompt
|
56 |
+
negative_prompt = Negative_prompt
|
57 |
+
n_steps = Steps
|
58 |
+
# run both experts
|
59 |
+
image = base(
|
60 |
+
prompt=prompt,
|
61 |
+
negative_prompt=negative_prompt,
|
62 |
+
num_inference_steps=n_steps,
|
63 |
+
denoising_end=high_noise_frac,
|
64 |
+
output_type="latent",
|
65 |
+
).images
|
66 |
+
image = refiner(
|
67 |
+
prompt=prompt,
|
68 |
+
num_inference_steps=n_steps,
|
69 |
+
denoising_start=high_noise_frac,
|
70 |
+
image=image,
|
71 |
+
).images[0]
|
72 |
+
return image
|
73 |
+
|
74 |
+
def predict(input_image):
|
75 |
+
model = torch.hub.load('RF5/danbooru-pretrained', 'resnet50')
|
76 |
+
model.eval()
|
77 |
+
|
78 |
+
# Load JSON file from github as Label
|
79 |
+
with urllib.request.urlopen("https://github.com/RF5/danbooru-pretrained/raw/master/config/class_names_6000.json") as url:
|
80 |
+
labels = json.loads(url.read().decode())
|
81 |
+
#Convert input image from array to PIL Image
|
82 |
+
input_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
|
83 |
+
#Preprocess the input image
|
84 |
+
preprocess = transforms.Compose([
|
85 |
+
transforms.Resize(360),
|
86 |
+
transforms.ToTensor(),
|
87 |
+
transforms.Normalize(mean=[0.7137, 0.6628, 0.6519], std=[0.2970, 0.3017, 0.2979]),
|
88 |
+
])
|
89 |
+
input_tensor = preprocess(input_image)
|
90 |
+
input_batch = input_tensor.unsqueeze(0)
|
91 |
+
|
92 |
+
# Use CUDA if available
|
93 |
+
if torch.cuda.is_available():
|
94 |
+
input_batch = input_batch.to('cuda')
|
95 |
+
model.to('cuda')
|
96 |
+
|
97 |
+
# Make prediction
|
98 |
+
with torch.no_grad():
|
99 |
+
output = model(input_batch)
|
100 |
+
|
101 |
+
# Get probabilities
|
102 |
+
probs = torch.sigmoid(output[0])
|
103 |
+
|
104 |
+
# Convert tensor to Python list of floats
|
105 |
+
probs = probs.cpu().numpy().tolist()
|
106 |
+
|
107 |
+
# Sort labels with probabilities and return top 10
|
108 |
+
sorted_labels_with_probs = sorted(list(zip(labels, probs)), key=lambda x: x[1], reverse=True)[:10]
|
109 |
+
|
110 |
+
# Convert list of tuples to dictionary and convert numpy floats to Python floats
|
111 |
+
sorted_labels_with_probs_dict = {label: float(prob) for label, prob in sorted_labels_with_probs}
|
112 |
+
return sorted_labels_with_probs_dict
|
113 |
+
|
114 |
+
def image_classify(input_image, model):
|
115 |
+
model_mapping= {
|
116 |
+
"Resnet 50": "microsoft/resnet-50",
|
117 |
+
"Vit Base Patch16-224": "google/vit-base-patch16-224",
|
118 |
+
"NSFW Image Detection": "Falconsai/nsfw_image_detection",
|
119 |
+
"Vit Age Classifier": "nateraw/vit-age-classifier"
|
120 |
+
}
|
121 |
+
classifier = pipeline("image-classification", model=model_mapping[model])
|
122 |
+
img = input_image
|
123 |
+
result = classifier(img)
|
124 |
+
#Sort the perccentage confident from highest to lowest
|
125 |
+
highest_confidence_result = sorted(result, key=lambda x: x['score'], reverse=True)[0]
|
126 |
+
# Format the score as a percentage and combine it with the label
|
127 |
+
output = f"{highest_confidence_result['score']*100:.2f}% confident : {highest_confidence_result['label']}"
|
128 |
+
return output
|
129 |
+
|
130 |
+
# Define the demo tab
|
131 |
+
with gr.Blocks() as demo:
|
132 |
+
with gr.Tab("Demo"):
|
133 |
+
image_input = gr.Image(type='pil')
|
134 |
+
image_output = gr.Image()
|
135 |
+
demo_button = gr.Button("Generate")
|
136 |
+
|
137 |
+
with gr.Tab("Text2Image"):
|
138 |
+
SD_text_input = gr.Textbox(lines=5, label="Prompt")
|
139 |
+
SD_text2_input = gr.Textbox(lines=5, label="Negative Prompt")
|
140 |
+
Slider_input = gr.Slider(0, 100, label="Strength")
|
141 |
+
SD_output = gr.Image()
|
142 |
+
SD_button = gr.Button("Generate")
|
143 |
+
|
144 |
+
with gr.Tab("Text2Image"):
|
145 |
+
option_input = gr.Dropdown(["resnet50", "vit-base-patch16-224", "vit-age-classifier", "nsfw image classification"], label="Model")
|
146 |
+
t2i_input = gr.Image(label="Image")
|
147 |
+
t2i_output = gr.Image()
|
148 |
+
t2i_button = gr.Button("Generate")
|
149 |
+
|
150 |
+
demo_button.click(demo_tab, inputs=image_input, outputs=image_output)
|
151 |
+
SD_button.click(generate_image, inputs=[SD_text_input,SD_text2_input,Slider_input], outputs=SD_output)
|
152 |
+
t2i_button.click(image_classify, inputs=[option_input,t2i_input], outputs=t2i_button)
|
153 |
+
|
154 |
+
demo.launch(debug=True)
|