File size: 1,423 Bytes
aa17401 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
# Load Model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "runwayml/stable-diffusion-v1-5"
pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipeline.safety_checker = None
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
pipeline = pipeline.to(device)
# Image generation function
def generate_image(prompt):
negative_prompt = "blurry, distorted, bad anatomy, missing fingers, extra limbs, bad proportions, deformed hands, low quality, cropped, out of frame, partial body, cut off, head cut off"
refined_prompt = f"{prompt}, centered, full body, well-framed, symmetrical"
with torch.no_grad():
image = pipeline(refined_prompt, negative_prompt=negative_prompt, guidance_scale=7.5, num_inference_steps=50).images[0]
torch.cuda.empty_cache() # Free GPU memory
return image
# Gradio UI
interface = gr.Interface(
fn=generate_image,
inputs=gr.Textbox(label="Enter Your Prompt"),
outputs=gr.Image(label="Generated Image"),
title="Rotona: Unrestricted Image Generator",
description="Enter a prompt to generate an AI image. This model is unrestricted—use responsibly."
)
# Launch
if __name__ == "__main__":
interface.launch()
|