import gradio as gr from share_btn import community_icon_html, loading_icon_html, share_js import random import re import torch from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline, set_seed from optimum.intel.openvino import OVStableDiffusionPipeline horoscope_model_id = "shahp7575/gpt2-horoscopes" tokenizer = AutoTokenizer.from_pretrained(horoscope_model_id) model = AutoModelWithLMHead.from_pretrained(horoscope_model_id) text_generation_pipe = pipeline("text-generation", model="Gustavosta/MagicPrompt-Stable-Diffusion", tokenizer="gpt2") stable_diffusion_pipe = OVStableDiffusionPipeline.from_pretrained("echarlaix/stable-diffusion-v1-5-openvino", revision="fp16", compile=False) height = 128 width = 128 stable_diffusion_pipe.reshape(batch_size=1, height=height, width=width, num_images_per_prompt=1) stable_diffusion_pipe.compile() def fn(sign, cat): prompt = f"<|category|> {cat} <|horoscope|> {sign}" prompt_encoded = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0) sample_outputs = model.generate( prompt_encoded, do_sample=True, top_k=40, max_length=300, top_p=0.95, temperature=0.95, num_beams=4, num_return_sequences=1, ) final_out = tokenizer.decode(sample_outputs[0], skip_special_tokens=True) starting_text = " ".join(final_out.split(" ")[4:]) seed = random.randint(100, 1000000) set_seed(seed) response = text_generation_pipe(starting_text + " " + sign + " art", max_length=(len(starting_text) + random.randint(60, 90)), num_return_sequences=1) image = stable_diffusion_pipe(response[0]["generated_text"], height=height, width=width, num_inference_steps=30).images[0] return [image, starting_text] block = gr.Blocks(css="./css.css") with block: with gr.Group(): with gr.Box(): with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True): text = gr.Dropdown( label="Star Sign", choices=["Aries", "Taurus","Gemini", "Cancer", "Leo", "Virgo", "Libra", "Scorpio", "Sagittarius", "Capricorn", "Aquarius", "Pisces"], show_label=True, max_lines=1, placeholder="Enter your prompt", elem_id="prompt-text-input", ).style( border=(True, False, True, True), rounded=(True, False, False, True), container=False, ) text2 = gr.Dropdown( choices=["Love", "Career", "Wellness"], label="Category", show_label=True, max_lines=1, placeholder="Enter your prompt", elem_id="prompt-text-input", ).style( border=(True, True, True, True), rounded=(True, False, False, True), container=False, ) btn = gr.Button("Generate image").style( margin=False, rounded=(False, True, True, False), full_width=False, ) gallery = gr.Image( interactive=False, label="Generated images", show_label=False, elem_id="gallery" ).style(grid=[2], height="auto") text = gr.Textbox("Text") with gr.Group(elem_id="container-advanced-btns"): with gr.Group(elem_id="share-btn-container"): community_icon = gr.HTML(community_icon_html) loading_icon = gr.HTML(loading_icon_html) share_button = gr.Button("Share to community", elem_id="share-btn") btn.click(fn=fn, inputs=[text, text2], outputs=[gallery, text]) share_button.click( None, [], [], _js=share_js, ) block.queue(concurrency_count=40, max_size=20).launch(max_threads=150)