Spaces:
Runtime error
Runtime error
| import torch | |
| from diffusers import StableDiffusion3Pipeline | |
| import gradio as gr | |
| import os | |
| import transformers | |
| import numpy as np | |
| from transformers import T5Tokenizer, T5ForConditionalGeneration | |
| import spaces | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| if torch.cuda.is_available(): | |
| device = "cuda" | |
| print("Using GPU") | |
| else: | |
| device = "cpu" | |
| print("Using CPU") | |
| MAX_SEED = np.iinfo(np.int32).max | |
| # Initialize the pipeline and download the sd3 medium model | |
| pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16) | |
| pipe.to(device) | |
| # superprompt-v1 | |
| tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1") | |
| model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype="auto") | |
| model.to(device) | |
| # Define the image generation function | |
| def generate_image(prompt, enhance_prompt, negative_prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt): | |
| if seed == 0: | |
| seed = random.randint(1, 2**32-1) | |
| if enhance_prompt: | |
| transformers.set_seed(seed) | |
| input_text = f"Expand the following prompt to add more detail: {prompt}" | |
| input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device) | |
| outputs = model.generate( | |
| input_ids, | |
| max_new_tokens=512, | |
| repetition_penalty=1.2, | |
| do_sample=True, | |
| temperature=0.7, | |
| top_p=1, | |
| top_k=50 | |
| ) | |
| prompt = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| generator = torch.Generator().manual_seed(seed) | |
| output = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| num_inference_steps=num_inference_steps, | |
| height=height, | |
| width=width, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| num_images_per_prompt=num_images_per_prompt | |
| ).images | |
| return output | |
| # Create the Gradio interface | |
| prompt = gr.Textbox(label="Prompt", info="Describe the image you want", placeholder="A cat...") | |
| enhance_prompt = gr.Checkbox(label="Prompt Enhancement", info="Enhance your prompt with SuperPrompt-v1", value=True) | |
| negative_prompt = gr.Textbox(label="Negative Prompt", info="Describe what you don't want in the image", value="deformed, distorted, disfigured, poorly drawn, bad anatomy, incorrect anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation", placeholder="Ugly, bad anatomy...") | |
| num_inference_steps = gr.Slider(label="Number of Inference Steps", info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference", minimum=1, maximum=50, value=25, step=1) | |
| height = gr.Slider(label="Height", info="Height of the Image", minimum=256, maximum=1344, step=32, value=1024) | |
| width = gr.Slider(label="Width", info="Width of the Image", minimum=256, maximum=1344, step=32, value=1024) | |
| guidance_scale = gr.Slider(label="Guidance Scale", info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.", minimum=0.0, maximum=10.0, value=7.5, step=0.1) | |
| seed = gr.Slider(value=42, minimum=0, maximum=MAX_SEED, step=1, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one") | |
| num_images_per_prompt = gr.Slider(label="Number of Images to generate with the settings",minimum=1, maximum=4, step=1, value=1) | |
| interface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[prompt, enhance_prompt, negative_prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt], | |
| outputs=gr.Gallery(label="Generated AI Images", elem_id="gallery"), | |
| title="Stable Diffusion 3 Medium", | |
| description="Made by <a href='https://linktr.ee/Nick088' target='_blank'>Nick088</a> \n Join https://discord.gg/osai to talk about Open Source AI" | |
| ) | |
| # Launch the interface | |
| interface.launch(share = False) |