app.py

import gradio as gr
import torch
from torchvision import transforms
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

# Define the starter and end prompts
starter_prompt = "Create a detailed sketch of a person based on the following description: "
end_prompt = " The sketch should have a monochromatic style, resembling a police mugshot, with clear outlines and shading to emphasize facial features. The background should be plain to focus on the subject."

# Function to generate the sketch
def generate_sketch(main_prompt):
    prompt = starter_prompt + main_prompt + end_prompt
    generator = torch.Generator("cpu").manual_seed(17)  # Use CPU for generation
    image = pipe(prompt, num_inference_steps=60, generator=generator).images[0]
    
    # Save the generated image (optional)
    image.save("testing.png")
    
    return image

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# InkScribe")
    main_prompt_input = gr.Textbox(label="Enter Description", placeholder="e.g., A young woman with curly black hair and green eyes, looking directly at the viewer.")
    generate_button = gr.Button("Generate Sketch")
    output_image = gr.Image(label="Generated Sketch", width= 512, height= 512)

    # Set up event handling
    generate_button.click(generate_sketch, inputs=main_prompt_input, outputs=output_image)

# Launch the app
demo.launch()