Spaces:
Running
Running
| import gradio as gr | |
| from google import genai | |
| from google.genai import types | |
| from PIL import Image | |
| from io import BytesIO | |
| import tempfile | |
| import concurrent.futures | |
| import os | |
| api_key = os.environ.get("API_KEY") | |
| univin_model = os.environ.get("univin") | |
| client = genai.Client(api_key=api_key) | |
| PROMPT_VARIATIONS = [ | |
| "Chain-of-Thought: Step 1: Carefully analyze the input image and note every detail of the product including its text. Step 2: Generate a high-quality image with a minimalist white background and bright natural light while keeping the product and all text exactly as in the original image. Step 3: Verify that the generated image is perfect and upload-ready.", | |
| "Chain-of-Thought: Step 1: Examine the product in the input image with focus on preserving every detail and text. Step 2: Generate a high-quality image featuring a beautiful woman using the product in a natural setting, ensuring the product and its text remain exactly unchanged. Step 3: Confirm that the final image meets professional standards and is ready for upload.", | |
| "Chain-of-Thought: Step 1: Inspect the input image thoroughly, capturing all aspects of the product and its text. Step 2: Generate a high-quality image from a slightly elevated angle in a modern indoor setting with warm lighting, while keeping the product and its text intact. Step 3: Validate that the resulting image is flawless and upload-ready.", | |
| "Chain-of-Thought: Step 1: Analyze the original product image and note all details, ensuring the text is preserved. Step 2: Generate a high-quality image using a low angle shot with cool ambient lighting and a contemporary studio background, without altering any aspect of the product or text. Step 3: Ensure the image is perfect and ready to be uploaded.", | |
| "Chain-of-Thought: Step 1: Study the input image closely, ensuring every detail of the product and its text is noted. Step 2: Generate a high-quality image with a vibrant outdoor background and natural sunlight, while keeping the product and its text completely unchanged. Step 3: Double-check that the final image is flawless and upload-ready.", | |
| "Chain-of-Thought: Step 1: Review the product image carefully, preserving all details including the text. Step 2: Generate a high-quality image with a subtle gradient background and soft diffused lighting, ensuring that no changes are made to the product or its text. Step 3: Confirm that the generated image is perfect and ready for upload.", | |
| "Chain-of-Thought: Step 1: Examine the original product image and identify every detail and text element. Step 2: Generate a high-quality image featuring an artistic abstract background with high contrast lighting, strictly preserving the product and its text. Step 3: Verify that the output image is impeccable and upload-ready.", | |
| "Chain-of-Thought: Step 1: Analyze the provided product image, capturing all details including text. Step 2: Generate a high-quality image with a dynamic angle and a colorful background, while ensuring the product and its text are exactly as in the input. Step 3: Ensure the final image is flawless and ready to be uploaded.", | |
| "Chain-of-Thought: Step 1: Thoroughly inspect the input image to capture every detail of the product and its text. Step 2: Generate a high-quality image with a scenic nature background and golden hour lighting, keeping the product and its text intact. Step 3: Confirm that the generated image is perfect and upload-ready.", | |
| "Chain-of-Thought: Step 1: Review the original product image with a focus on preserving all details and text. Step 2: Generate a high-quality image where a beautiful woman is using the product, ensuring that the product and its text remain exactly unchanged. Step 3: Validate that the final image meets perfection and is ready for upload." | |
| ] | |
| def process_variation(variation, input_image, product_name): | |
| text_input = ( | |
| f"Hi, this is a picture of a product. The name of the product is {product_name}.", | |
| variation | |
| ) | |
| response = client.models.generate_content( | |
| model=univin_model, | |
| contents=[text_input, input_image], | |
| config=types.GenerateContentConfig(response_modalities=['Text', 'Image']) | |
| ) | |
| for part in response.candidates[0].content.parts: | |
| if part.inline_data is not None: | |
| generated_img = Image.open(BytesIO(part.inline_data.data)) | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file: | |
| generated_img.save(tmp_file, format="PNG") | |
| return tmp_file.name | |
| return None | |
| def generate_images(input_image, product_name): | |
| with concurrent.futures.ThreadPoolExecutor() as executor: | |
| futures = [ | |
| executor.submit(process_variation, variation, input_image, product_name) | |
| for variation in PROMPT_VARIATIONS | |
| ] | |
| output_files = [future.result() for future in futures if future.result() is not None] | |
| return output_files | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Uni-Imaginator") | |
| with gr.Row(): | |
| input_image = gr.Image(type="pil", label="Upload Image") | |
| product_name = gr.Textbox(label="Product Name", placeholder="Enter the product name") | |
| generate_button = gr.Button("Generate Images") | |
| gallery = gr.Gallery(label="Generated Images", elem_id="gallery", columns=4, height="auto", object_fit="contain", interactive=True) | |
| generate_button.click(fn=generate_images, inputs=[input_image, product_name], outputs=gallery) | |
| demo.launch() | |