Spaces:
Build error
Build error
| import torch | |
| from diffusers import AutoPipelineForImage2Image | |
| from PIL import Image, ImageDraw, ImageFont | |
| import requests | |
| from io import BytesIO | |
| import gradio as gr | |
| # log gpu availabilitu | |
| print(f"Is CUDA available: {torch.cuda.is_available()}") | |
| print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}") | |
| def image_to_template(generated_image, logo, button_text, punchline, theme_color): | |
| # Resize logo if needed | |
| logo = logo.resize((100, 100)) # Example size, adjust as needed | |
| # Create a blank canvas with extra space for logo, punchline, and button | |
| canvas_width = max(generated_image.width, logo.width) * 2 | |
| canvas_height = generated_image.height + logo.height + 100 | |
| canvas = Image.new('RGB', (canvas_width, canvas_height), 'white') | |
| # Paste the logo and the generated image onto the canvas | |
| canvas.paste(logo, (10, 10)) # Adjust position as needed | |
| canvas.paste(generated_image, (0, logo.height + 20)) | |
| # Add punchline and button | |
| draw = ImageDraw.Draw(canvas) | |
| font = ImageFont.load_default() # Or use a custom font | |
| text_color = theme_color | |
| # Draw punchline | |
| draw.text((10, logo.height + generated_image.height + 30), punchline, fill=text_color, font=font) | |
| # Draw button | |
| button_position = (10, logo.height + generated_image.height + 60) # Adjust as needed | |
| draw.rectangle([button_position, (canvas_width - 10, canvas_height - 10)], outline=theme_color, fill=text_color) | |
| draw.text(button_position, button_text, font=font) | |
| return canvas | |
| def generate_template(initial_image, logo, prompt, button_text, punchline, image_color, theme_color): | |
| pipeline = AutoPipelineForImage2Image.from_pretrained( | |
| "./models/kandinsky-2-2-decoder", | |
| torch_dtype=torch.float16, | |
| variant="fp16", | |
| use_safetensors=True | |
| ) | |
| # pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True) | |
| pipeline.enable_model_cpu_offload() | |
| prompt = f"{prompt}, include the color {image_color}" | |
| negative_prompt = "low quality, bad quality, blurry, unprofessional" | |
| generated_image = pipeline( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| image=initial_image, | |
| height=256, | |
| width=256).images[0] | |
| template_image = image_to_template(generated_image, logo, button_text, punchline, theme_color) | |
| return template_image | |
| # Set up Gradio interface | |
| iface = gr.Interface( | |
| fn=generate_template, | |
| inputs=[gr.inputs.Image(type="pil", label="Initial Image"), | |
| gr.inputs.Image(type="pil", label="Logo"), | |
| gr.inputs.Textbox(label="Prompt"), | |
| gr.inputs.Textbox(label="Button Text"), | |
| gr.inputs.Textbox(label="Punchline"), | |
| gr.inputs.ColorPicker(label="Image Color"), | |
| gr.inputs.ColorPicker(label="Theme Color")], | |
| outputs=[gr.outputs.Image(type="pil")], | |
| title="Ad Template Generation Using Diffusion Models Demo", | |
| description="Generate ad template based on your inputs using a trained model.", | |
| concurrency_limit=2, | |
| # examples=[ | |
| # [] | |
| # ] | |
| ) | |
| # Run the interface | |
| iface.launch() | |