PlayForLose commited on
Commit
19aec01
1 Parent(s): 70138aa

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +154 -0
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)