Spaces:
Runtime error
Runtime error
| import json #to load json file for custom words for profanity check | |
| import torch | |
| import spaces | |
| from PIL import Image | |
| from generationPipeline import generate | |
| from transformers import CLIPTokenizer | |
| from loadModel import preload_models_from_standard_weights | |
| from better_profanity import profanity # Import the profanity-check library | |
| import gradio as gr | |
| Device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| print(f"Using device: {Device}") | |
| tokenizer = CLIPTokenizer("vocab.json", merges_file="merges.txt") | |
| model_file = "weights-inkpen.ckpt" | |
| models = preload_models_from_standard_weights(model_file, Device) | |
| ## profanity check on input prompt | |
| profanity.load_censor_words() | |
| ##fFunction to load custom words from json file | |
| def load_custom_words(file_path): | |
| with open(file_path, 'r') as file: | |
| data = json.load(file) | |
| return data["custom_words"] | |
| #3 load and add custom words from the json file | |
| custom_words = load_custom_words("profane.json") | |
| profanity.add_censor_words(custom_words) | |
| def filter_prompt(prompt): | |
| if profanity.contains_profanity(prompt): | |
| return "Inappropriate content detected. Please modify the input." | |
| return prompt | |
| def generate_image(mode, prompt, strength, seed, n_inference_steps, input_image=None): | |
| ##check prompt id there is anything inappropriate | |
| filtered_prompt = filter_prompt(prompt) | |
| if filtered_prompt == "Inappropriate content detected. Please modify the input.": | |
| return filtered_prompt | |
| if mode == "Text-to-Image": | |
| # Ignore the input image | |
| output_image = generate( | |
| prompt=filtered_prompt, | |
| uncond_prompt="", | |
| input_image=None, | |
| strength=strength, | |
| do_cfg=True, | |
| cfg_scale=8, | |
| sampler_name="ddpm", | |
| n_inference_steps=n_inference_steps, | |
| seed=seed, | |
| models=models, | |
| device=Device, | |
| idle_device="cpu", | |
| tokenizer=tokenizer, | |
| ) | |
| elif mode == "Image-to-Image" and input_image is not None: | |
| # Use the uploaded image | |
| output_image = generate( | |
| prompt=filtered_prompt, | |
| uncond_prompt="", | |
| input_image=input_image, | |
| strength=strength, | |
| do_cfg=True, | |
| cfg_scale=8, | |
| sampler_name="ddpm", | |
| n_inference_steps=n_inference_steps, | |
| seed=seed, | |
| models=models, | |
| device=Device, | |
| idle_device="cpu", | |
| tokenizer=tokenizer, | |
| ) | |
| else: | |
| return "Please upload an image for Image-to-Image mode." | |
| return Image.fromarray(output_image) | |
| # Gradio interface with inputs for mode selection and inference steps | |
| iface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.Radio(choices=["Text-to-Image", "Image-to-Image"], label="Mode"), # Toggle between modes | |
| gr.Textbox(label="Text Prompt"), # Input text prompt | |
| gr.Slider(0, 1, step=0.01, label="Strength : (Note: set strength between 0.01 to 0.9, For Image-2-Image: Strength ~ 1 means that the output will be further from the input image. Strength ~ 0 means that the output will be closer to the input image.)"), # Slider for strength | |
| gr.Number(label="Seed (for reproducibility)"), # Seed for reproducibility | |
| gr.Slider(1, 200, step=10, label="Number of Inference Steps"), # Slider for inference steps | |
| gr.Image(type='filepath', label='Input Image Only for Image-to-Image'), # Input image without optional flag | |
| ], | |
| outputs=gr.Image(label="Generated Image"), # Output generated image | |
| title="Stable Diffusion Image Generator", | |
| description="Generate images from text prompts or use image-to-image generation." | |
| ) | |
| # Launch the Gradio interface with public sharing enabled | |
| iface.launch(share=True, debug=True) |