Spaces:
Running
Running
File size: 4,234 Bytes
7139c6f |
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 99 100 101 102 103 104 105 106 107 108 109 110 |
import gradio as gr
import requests
import io
import os
import logging
from PIL import Image, ImageEnhance
from gradio_client import Client
# API URLs and headers
API_URL_1 = "https://api-inference.huggingface.co/models/sd-community/sdxl-flash"
API_KEY = os.getenv("HF_API_KEY")
headers = {"Authorization": f"Bearer {API_KEY}"}
client = Client("AP123/SDXL-Lightning")
def query(api_url, payload):
try:
response = requests.post(api_url, headers=headers, json=payload)
response.raise_for_status()
return response.content
except requests.exceptions.HTTPError as http_err:
logging.error(f"HTTP error occurred: {http_err}")
print(f"HTTP error occurred: {http_err}")
except Exception as err:
logging.error(f"An error occurred: {err}")
print(f"An error occurred: {err}")
return None
def map_inference_steps(steps):
if steps <= 1:
return "1-Step"
elif steps <= 2:
return "2-Step"
elif steps <= 4:
return "4-Step"
else:
return "8-Step"
def gen_img(prompt, num_inference_steps, guidance_scale, contrast):
payload = {
"inputs": prompt,
"parameters": {
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale
}
}
# Try generating image with the first model
image_bytes = query(API_URL_1, payload)
if image_bytes is None:
# print("First model failed, trying second model...")
try:
mapped_steps = map_inference_steps(num_inference_steps)
result = client.predict(prompt, mapped_steps, api_name="/generate_image")
with open(result, "rb") as img_file:
image = Image.open(img_file)
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(contrast)
return image, "Image generated by the second model"
except Exception as e:
# logging.error(f"Error with the second model: {e}")
# print(f"Error with the second model: {e}")
return None, "Error with the second model"
try:
image = Image.open(io.BytesIO(image_bytes))
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(contrast)
return image, "Image generated by the first model"
except Exception as e:
# logging.error(f"Error processing image: {e}")
return None, "Error processing image"
with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.lime, spacing_size="md", radius_size="lg",
font=[gr.themes.GoogleFont("Copperplate"), "Fantasy", "sans-serif"])) as demo:
gr.Markdown("""
<div style='text-align: center;
font-size: 36px;
font-family: Copperplate, Fantasy;
color: orange;'>
Pencil Vision (early beta 1.3)
</div>
""")
with gr.Row():
prompt = gr.Textbox(label="Enter your prompt to generate the image", placeholder="A fantasy landscape with mountains and a river")
with gr.Row():
num_inference_steps = gr.Slider(minimum=1, maximum=15, value=6, step=1, label="Number of Inference Steps")
guidance_scale = gr.Slider(minimum=1.0, maximum=5.0, value=2.5, step=0.1, label="Guidance Scale")
contrast_scale = gr.Slider(minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Contrast Scale")
with gr.Row():
generate_button = gr.Button("Generate Image")
output_image = gr.Image(label="Generated Image")
# model_info = gr.Textbox(label="Model Info", interactive=False)
error_message = gr.Markdown(visible=False)
def handle_click(prompt, num_inference_steps, guidance_scale, contrast_scale):
image, info = gen_img(prompt, num_inference_steps, guidance_scale, contrast_scale)
if image is None:
return None, info, gr.update(visible=True, value="Error generating image. Please try again.")
return image, info, gr.update(visible=True)
generate_button.click(handle_click, inputs=[prompt, num_inference_steps, guidance_scale, contrast_scale], outputs=[output_image, error_message])
demo.launch(share=True)
|