Upload 5 files

app.py

@@ -0,0 +1,169 @@
+import gradio as gr
+from PIL import Image
+import asyncio
+
+from text_to_image import TextToImage
+from image_to_text import ImageToText
+from image_to_image import ImageToImage
+
+# ============================================
+# Initialize Model Classes
+# ============================================
+
+text_to_image = TextToImage()
+image_to_text = ImageToText()
+image_to_image = ImageToImage()
+
+# ============================================
+# Gradio Interface Functions with Async and Error Handling
+# ============================================
+
+async def async_text_to_image(prompt):
+    """
+    Asynchronous interface function for Text-to-Image generation with error handling.
+    """
+    try:
+        image = await text_to_image.generate_image(prompt)
+        return image
+    except Exception as e:
+        raise gr.Error(f"Text-to-Image Generation Failed: {str(e)}")
+
+async def async_image_to_text(image):
+    """
+    Asynchronous interface function for Image-to-Text captioning with error handling.
+    """
+    try:
+        caption = await image_to_text.generate_caption(image)
+        return caption
+    except Exception as e:
+        raise gr.Error(f"Image-to-Text Captioning Failed: {str(e)}")
+
+async def async_image_to_image(image, prompt):
+    """
+    Asynchronous interface function for Image-to-Image transformation with error handling.
+    """
+    try:
+        transformed_image = await image_to_image.transform_image(image, prompt)
+        return transformed_image
+    except Exception as e:
+        raise gr.Error(f"Image-to-Image Transformation Failed: {str(e)}")
+
+# ============================================
+# Gradio UI Design
+# ============================================
+
+with gr.Blocks(css=".gradio-container {background-color: #f0f8ff}") as demo:
+    # Title Section
+    gr.Markdown("# 🎨 AI Creativity Hub 🚀")
+    gr.Markdown("### Unleash the power of AI to transform your ideas into reality!")
+
+    # Task Selection Radio
+    with gr.Tab("✨ Choose Your Magic ✨"):
+        task = gr.Radio(
+            ["🖼️ Text-to-Image", "📝 Image-to-Text", "🖌️ Image-to-Image"],
+            label="Select a Task",
+            interactive=True,
+            value="🖼️ Text-to-Image"
+        )
+
+    # Text-to-Image Section
+    with gr.Row(visible=False) as text_to_image_tab:
+        with gr.Column():
+            gr.Markdown("## 🖼️ Text-to-Image Generator")
+            prompt_input = gr.Textbox(
+                label="📝 Enter your prompt:",
+                placeholder="e.g., A serene sunset over the mountains",
+                lines=2
+            )
+            generate_btn = gr.Button("🎨 Generate Image")
+            with gr.Row():
+                output_image = gr.Image(label="🖼️ Generated Image")
+                download_btn = gr.Button("📥 Download Image")
+    
+    # Image-to-Text Section
+    with gr.Row(visible=False) as image_to_text_tab:
+        with gr.Column():
+            gr.Markdown("## 📝 Image-to-Text Captioning")
+            image_input = gr.Image(
+                label="📸 Upload an image:",
+                type="pil"
+            )
+            generate_caption_btn = gr.Button("🖋️ Generate Caption")
+            caption_output = gr.Textbox(
+                label="📝 Generated Caption:",
+                lines=2
+            )
+    
+    # Image-to-Image Section
+    with gr.Row(visible=False) as image_to_image_tab:
+        with gr.Column():
+            gr.Markdown("## 🖌️ Image-to-Image Transformer")
+            init_image_input = gr.Image(
+                label="📸 Upload an image:",
+                type="pil"
+            )
+            transformation_prompt = gr.Textbox(
+                label="📝 Enter transformation prompt:",
+                placeholder="e.g., Make it look like a Van Gogh painting",
+                lines=2
+            )
+            transform_btn = gr.Button("🔄 Transform Image")
+            with gr.Row():
+                transformed_image = gr.Image(label="🖌️ Transformed Image")
+                download_transformed_btn = gr.Button("📥 Download Image")
+    
+    # Define Visibility Based on Task Selection
+    def toggle_visibility(selected_task):
+        return {
+            text_to_image_tab: selected_task == "🖼️ Text-to-Image",
+            image_to_text_tab: selected_task == "📝 Image-to-Text",
+            image_to_image_tab: selected_task == "🖌️ Image-to-Image",
+        }
+    
+    task.change(
+        fn=toggle_visibility, 
+        inputs=task, 
+        outputs=[text_to_image_tab, image_to_text_tab, image_to_image_tab]
+    )
+
+    # Define Button Actions
+    generate_btn.click(
+        fn=async_text_to_image, 
+        inputs=prompt_input, 
+        outputs=output_image
+    )
+    
+    download_btn.click(
+        fn=lambda img: img.save("generated_image.png") or "Image downloaded!",
+        inputs=output_image, 
+        outputs=None
+    )
+    
+    generate_caption_btn.click(
+        fn=async_image_to_text, 
+        inputs=image_input, 
+        outputs=caption_output
+    )
+    
+    transform_btn.click(
+        fn=async_image_to_image, 
+        inputs=[init_image_input, transformation_prompt], 
+        outputs=transformed_image
+    )
+    
+    download_transformed_btn.click(
+        fn=lambda img: img.save("transformed_image.png") or "Image downloaded!",
+        inputs=transformed_image, 
+        outputs=None
+    )
+    
+    # Footer Section with Quirky Elements
+    gr.Markdown("----")
+    gr.Markdown("### 🌟 Explore the endless possibilities with AI! 🌟")
+    gr.Markdown("#### 🚀 Built with ❤️ by [Your Name]")
+
+# ============================================
+# Launch the Gradio App
+# ============================================
+
+demo.launch()

image_to_image.py

@@ -0,0 +1,60 @@
+import torch
+from diffusers import StableDiffusionXLImg2ImgPipeline, EulerDiscreteScheduler
+from PIL import Image
+from io import BytesIO
+from utils import load_unet_model
+
+class ImageToImage:
+    """
+    Class to handle Image-to-Image transformations using Stable Diffusion XL.
+    """
+    def __init__(self, device="cpu"):
+        # Model and repository details
+        self.base = "stabilityai/stable-diffusion-xl-base-1.0"
+        self.repo = "ByteDance/SDXL-Lightning"
+        self.ckpt = "sdxl_lightning_4step_unet.safetensors"
+        self.device = device
+
+        # Load the UNet model
+        print("Loading Image-to-Image model...")
+        self.unet = load_unet_model(self.base, self.repo, self.ckpt, device=self.device)
+        
+        # Initialize the pipeline
+        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            self.base, 
+            unet=self.unet, 
+            torch_dtype=torch.float32
+        ).to(self.device)
+        
+        # Set the scheduler
+        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
+            self.pipe.scheduler.config, 
+            timestep_spacing="trailing"
+        )
+        print("Image-to-Image model loaded successfully.")
+    
+    
+    async def transform_image(self, image, prompt):
+        """
+        Transform an uploaded image based on a text prompt.
+
+        Args:
+            image (PIL.Image): The input image to transform.
+            prompt (str): The text prompt to guide the transformation.
+
+        Returns:
+            PIL.Image: The transformed image.
+        """
+        if not prompt:
+            raise ValueError("Prompt cannot be empty.")
+
+        # Resize the image as required by the model
+        init_image = image.resize((768, 512))
+        with torch.no_grad():
+            transformed_image = self.pipe(
+                prompt=prompt, 
+                image=init_image, 
+                strength=0.75, 
+                guidance_scale=7.5
+            ).images[0]
+        return transformed_image

image_to_text.py

@@ -0,0 +1,29 @@
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from PIL import Image
+
+class ImageToText:
+    """
+    Class to handle Image-to-Text captioning using BLIP.
+    """
+    def __init__(self):
+        # Initialize the processor and model
+        print("Loading Image-to-Text model...")
+        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+        self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+        print("Image-to-Text model loaded successfully.")
+    
+    
+    async def generate_caption(self, image):
+        """
+        Generate a descriptive caption for an uploaded image.
+
+        Args:
+            image (PIL.Image): The image to caption.
+
+        Returns:
+            str: The generated caption.
+        """
+        inputs = self.processor(image, return_tensors="pt")
+        out = self.model.generate(**inputs)
+        caption = self.processor.decode(out[0], skip_special_tokens=True)
+        return caption

text_to_image.py

@@ -0,0 +1,55 @@
+import torch
+import spaces
+from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from PIL import Image
+from io import BytesIO
+from utils import load_unet_model
+
+@spaces.GPU
+class TextToImage:
+    """
+    Class to handle Text-to-Image generation using Stable Diffusion XL.
+    """
+    def __init__(self, device="cpu"):
+        # Model and repository details
+        self.base = "stabilityai/stable-diffusion-xl-base-1.0"
+        self.repo = "ByteDance/SDXL-Lightning"
+        self.ckpt = "sdxl_lightning_4step_unet.safetensors"
+        self.device = device
+
+        # Load the UNet model
+        print("Loading Text-to-Image model...")
+        self.unet = load_unet_model(self.base, self.repo, self.ckpt, device=self.device)
+        
+        # Initialize the pipeline
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            self.base, 
+            unet=self.unet, 
+            torch_dtype=torch.float32,
+        ).to(self.device)
+        
+        # Set the scheduler
+        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
+            self.pipe.scheduler.config, 
+            timestep_spacing="trailing"
+        )
+        print("Text-to-Image model loaded successfully.")
+    
+    
+    async def generate_image(self, prompt):
+        """
+        Generate an image from a text prompt.
+
+        Args:
+            prompt (str): The text prompt to generate the image.
+
+        Returns:
+            PIL.Image: The generated image.
+        """
+        with torch.no_grad():
+            image = self.pipe(
+                prompt, 
+                num_inference_steps=4, 
+                guidance_scale=0
+            ).images[0]
+        return image

utils.py

@@ -0,0 +1,21 @@
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+def load_unet_model(base, repo, ckpt, device="cpu"):
+    """
+    Load the UNet model from Hugging Face Hub.
+
+    Args:
+        base (str): Base model name.
+        repo (str): Repository name.
+        ckpt (str): Checkpoint filename.
+        device (str): Device to load the model on.
+
+    Returns:
+        UNet2DConditionModel: Loaded UNet model.
+    """
+    from diffusers import UNet2DConditionModel
+    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, torch.float16)
+    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
+    return unet