Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| """Untitled3.ipynb | |
| Automatically generated by Colaboratory. | |
| Original file is located at | |
| https://colab.research.google.com/drive/1BltKPv_n-glCuuIIYSBA6GHK-tmwbl20 | |
| """ | |
| !pip install fastai | |
| !pip install gradio --upgrade | |
| !pip install lida | |
| !pip install diffusers --upgrade | |
| !pip install tensorflow-probability --upgrade | |
| !pip install invisible_watermark transformers accelerate safetensors | |
| !pip install torch --upgrade | |
| !pip install transformers --upgrade | |
| !pip install datasets | |
| #!pip install typing_extensions==4.5.0 | |
| #!pip install fastapi==0.103.0 | |
| import torch | |
| from PIL import Image | |
| from torchvision import transforms | |
| import gradio as gr | |
| import json | |
| import urllib, urllib.request | |
| from diffusers import DiffusionPipeline | |
| from transformers import pipeline | |
| # First Page | |
| def demo_tab(image): | |
| # The demo tab simply returns the same image as input | |
| return image | |
| # Secibd Page | |
| def generate_image(Prompt,Negative_prompt,Steps): | |
| # load both base & refiner | |
| base = DiffusionPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True | |
| ) | |
| base.to("cuda") | |
| refiner = DiffusionPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-xl-refiner-1.0", | |
| text_encoder_2=base.text_encoder_2, | |
| vae=base.vae, | |
| torch_dtype=torch.float16, | |
| use_safetensors=True, | |
| variant="fp16", | |
| ) | |
| refiner.to("cuda") | |
| # Define how many steps and what % of steps to be run on each experts (80/20) here | |
| high_noise_frac = 0.8 | |
| prompt = Prompt | |
| negative_prompt = Negative_prompt | |
| n_steps = Steps | |
| # run both experts | |
| image = base( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| num_inference_steps=n_steps, | |
| denoising_end=high_noise_frac, | |
| output_type="latent", | |
| ).images | |
| image = refiner( | |
| prompt=prompt, | |
| num_inference_steps=n_steps, | |
| denoising_start=high_noise_frac, | |
| image=image, | |
| ).images[0] | |
| return image | |
| def predict(input_image): | |
| model = torch.hub.load('RF5/danbooru-pretrained', 'resnet50') | |
| model.eval() | |
| # Load JSON file from github as Label | |
| with urllib.request.urlopen("https://github.com/RF5/danbooru-pretrained/raw/master/config/class_names_6000.json") as url: | |
| labels = json.loads(url.read().decode()) | |
| #Convert input image from array to PIL Image | |
| input_image = Image.fromarray(input_image.astype('uint8'), 'RGB') | |
| #Preprocess the input image | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(360), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.7137, 0.6628, 0.6519], std=[0.2970, 0.3017, 0.2979]), | |
| ]) | |
| input_tensor = preprocess(input_image) | |
| input_batch = input_tensor.unsqueeze(0) | |
| # Use CUDA if available | |
| if torch.cuda.is_available(): | |
| input_batch = input_batch.to('cuda') | |
| model.to('cuda') | |
| # Make prediction | |
| with torch.no_grad(): | |
| output = model(input_batch) | |
| # Get probabilities | |
| probs = torch.sigmoid(output[0]) | |
| # Convert tensor to Python list of floats | |
| probs = probs.cpu().numpy().tolist() | |
| # Sort labels with probabilities and return top 10 | |
| sorted_labels_with_probs = sorted(list(zip(labels, probs)), key=lambda x: x[1], reverse=True)[:10] | |
| # Convert list of tuples to dictionary and convert numpy floats to Python floats | |
| sorted_labels_with_probs_dict = {label: float(prob) for label, prob in sorted_labels_with_probs} | |
| return sorted_labels_with_probs_dict | |
| def image_classify(input_image, model): | |
| model_mapping= { | |
| "Resnet 50": "microsoft/resnet-50", | |
| "Vit Base Patch16-224": "google/vit-base-patch16-224", | |
| "NSFW Image Detection": "Falconsai/nsfw_image_detection", | |
| "Vit Age Classifier": "nateraw/vit-age-classifier" | |
| } | |
| classifier = pipeline("image-classification", model=model_mapping[model]) | |
| img = input_image | |
| result = classifier(img) | |
| #Sort the perccentage confident from highest to lowest | |
| highest_confidence_result = sorted(result, key=lambda x: x['score'], reverse=True)[0] | |
| # Format the score as a percentage and combine it with the label | |
| output = f"{highest_confidence_result['score']*100:.2f}% confident : {highest_confidence_result['label']}" | |
| return output | |
| # Define the demo tab | |
| with gr.Blocks() as demo: | |
| with gr.Tab("Demo"): | |
| image_input = gr.Image(type='pil') | |
| image_output = gr.Image() | |
| demo_button = gr.Button("Generate") | |
| with gr.Tab("Text2Image"): | |
| SD_text_input = gr.Textbox(lines=5, label="Prompt") | |
| SD_text2_input = gr.Textbox(lines=5, label="Negative Prompt") | |
| Slider_input = gr.Slider(0, 100, label="Strength") | |
| SD_output = gr.Image() | |
| SD_button = gr.Button("Generate") | |
| with gr.Tab("Text2Image"): | |
| option_input = gr.Dropdown(["resnet50", "vit-base-patch16-224", "vit-age-classifier", "nsfw image classification"], label="Model") | |
| t2i_input = gr.Image(label="Image") | |
| t2i_output = gr.Image() | |
| t2i_button = gr.Button("Generate") | |
| demo_button.click(demo_tab, inputs=image_input, outputs=image_output) | |
| SD_button.click(generate_image, inputs=[SD_text_input,SD_text2_input,Slider_input], outputs=SD_output) | |
| t2i_button.click(image_classify, inputs=[option_input,t2i_input], outputs=t2i_button) | |
| demo.launch(debug=True) |