Spaces:
Runtime error
Runtime error
File size: 3,474 Bytes
84c9848 3d7d7db 84c9848 3d7d7db 84c9848 3d7d7db 84c9848 3d7d7db 84c9848 3d7d7db 84c9848 3d7d7db 84c9848 |
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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionImg2ImgPipeline
MODEL_ID = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
print(f"Loading model on device: {device}")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
low_cpu_mem_usage=True, # Added optimization for low RAM environments
)
pipe = pipe.to(device)
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {e}")
pipe = None
print("Image generation disabled due to model loading error.")
def generate_avatar(
initial_image: Image,
prompt: str,
style_strength: float,
steps: int,
negative_prompt: str
):
"""
Generates a new image based on an input image and a text prompt.
The 'style_strength' controls how much the new image deviates from the original.
"""
if pipe is None:
if not prompt:
return initial_image, "Please enter a descriptive text prompt."
if initial_image is None:
return initial_image, "Please upload a reference image."
print(f"Generating image with prompt: '{prompt}' and strength: {style_strength}")
try:
with torch.autocast(device) if device == "cuda" else torch.no_grad():
generated_image = pipe(
prompt=prompt,
image=initial_image.convert("RGB"), # Ensure image is in RGB format
strength=style_strength,
num_inference_steps=steps,
guidance_scale=7.5,
negative_prompt=negative_prompt
).images[0]
return generated_image, "Avatar successfully generated!"
except Exception as e:
print(f"Inference error: {e}")
return initial_image, f"An error occurred during generation: {e}"
# Define the input components
image_input = gr.Image(type="pil", label="1. Upload Base Image (Your Photo or Style Reference)")
prompt_input = gr.Textbox(
label="2. Creative Prompt (e.g., 'A professional LinkedIn headshot, cyberpunk aesthetic, highly detailed')",
placeholder="Describe the desired style and subject."
)
strength_slider = gr.Slider(
minimum=0.1,
maximum=0.9,
value=0.65,
step=0.05,
label="3. Transformation Strength (Higher = More stylistic change)"
)
steps_slider = gr.Slider(
minimum=10,
maximum=100,
value=30,
step=5,
label="4. Inference Steps (Higher = Better quality, Slower generation)"
)
negative_prompt_input = gr.Textbox(
label="5. Negative Prompt (What to AVOID, e.g., 'ugly, blurry, low quality, cropped')",
placeholder="Enter elements to avoid in the generated image."
)
gr.Interface(
fn=generate_avatar,
inputs=[image_input, prompt_input, strength_slider, steps_slider, negative_prompt_input],
outputs=[gr.Image(type="pil", label="Generated Profile Picture"), gr.Textbox(label="Status")],
title="✨ Creative Profile Picture Generator (SFW)",
description="Upload a reference photo and a creative prompt to generate a new, stylized avatar using image-to-image diffusion. Transformation strength controls how much the output deviates from the input image (0.1 is subtle, 0.9 is highly stylized).",
live=False,
allow_flagging="never"
).launch()
|