Spaces:
Runtime error
Runtime error
File size: 2,646 Bytes
f6b8b7e 9872917 c7a48c9 789e6b5 9872917 12c6b4e b85438c 789e6b5 66e43e3 13298a2 216a041 f466dd9 06a3d1b c7a48c9 9872917 789e6b5 c301a62 6449f8f f466dd9 b5ad13a 789e6b5 d26a101 789e6b5 9872917 e0ec116 b0bcf89 e0ec116 6449f8f 789e6b5 9872917 b0bcf89 06a3d1b 02161de 06a3d1b 789e6b5 06a3d1b f45116f c7a48c9 b0bcf89 6035350 f466dd9 9872917 f466dd9 789e6b5 d05fa5e 1adc78a e0ec116 d05fa5e c7a48c9 58f74fc f466dd9 c7a48c9 |
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 |
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, AutoPipelineForText2Image
from io import BytesIO
import asyncio
from generate_propmts import generate_prompt
from concurrent.futures import ThreadPoolExecutor, as_completed
# Load the model once outside of the function
model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
async def generate_image(prompt):
try:
# Generate an image based on the prompt
output = await asyncio.to_thread(model, prompt=prompt, num_inference_steps=50, guidance_scale=7.5)
print(f"Model output: {output}")
# Check if the model returned images
if isinstance(output.images, list) and len(output.images) > 0:
image = output.images[0]
buffered = BytesIO()
image.save(buffered, format="JPEG")
image_bytes = buffered.getvalue()
return image_bytes
else:
raise Exception("No images returned by the model.")
except Exception as e:
print(f"Error generating image: {e}")
return None
async def process_prompt(sentence_mapping, character_dict, selected_style):
images = {}
print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
prompts = []
# Generate prompts for each paragraph
for paragraph_number, sentences in sentence_mapping.items():
combined_sentence = " ".join(sentences)
prompt = generate_prompt(combined_sentence, sentence_mapping, character_dict, selected_style)
prompts.append((paragraph_number, prompt))
print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
# Create tasks for all prompts and run them concurrently
tasks = [generate_image(prompt) for _, prompt in prompts]
results = await asyncio.gather(*tasks)
# Map results back to paragraphs
for i, (paragraph_number, _) in enumerate(prompts):
if i < len(results):
images[paragraph_number] = results[i]
else:
print(f"Error: No image generated for paragraph {paragraph_number}")
return images
# Gradio interface with high concurrency limit
gradio_interface = gr.Interface(
fn=process_prompt,
inputs=[
gr.JSON(label="Sentence Mapping"),
gr.JSON(label="Character Dict"),
gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
],
outputs="json",
concurrency_count=20 # Set a high concurrency limit
)
if __name__ == "__main__":
gradio_interface.launch() # No need for share=True for local testing
|