img caption

app.py

@@ -1,138 +1,25 @@
 import gradio as gr
-import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
-from PIL import Image
-import io
-import base64
+from transformers import pipeline
 
-# Check if GPU is available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load image captioning model
+captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-# Load text-to-image pipeline
-pipe_text_to_image = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-).to(device)
-
-# Load image-to-image pipeline
-pipe_image_to_image = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-).to(device)
-
-def generate_image(prompt, negative_prompt="", steps=20, guidance=7.5):
-    """Generate image from text prompt"""
-    if not prompt.strip():
-        return None, "Please enter a prompt"
-    
-    try:
-        with torch.autocast(device):
-            image = pipe_text_to_image(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=steps,
-                guidance_scale=guidance
-            ).images[0]
-        return image, "Image generated successfully!"
-    except Exception as e:
-        return None, f"Error: {str(e)}"
-
-def modify_image(image, prompt, strength=0.75, steps=20, guidance=7.5):
-    """Modify existing image with prompt"""
+def caption_image(image):
     if image is None:
-        return None, "Please upload an image first"
-    if not prompt.strip():
-        return None, "Please enter a modification prompt"
+        return "Please upload an image"
     
-    try:
-        # Resize image to appropriate dimensions
-        width, height = image.size
-        max_size = 512
-        if width > max_size or height > max_size:
-            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-        
-        with torch.autocast(device):
-            result = pipe_image_to_image(
-                prompt=prompt,
-                image=image,
-                strength=strength,
-                num_inference_steps=steps,
-                guidance_scale=guidance
-            ).images[0]
-        return result, "Image modified successfully!"
-    except Exception as e:
-        return None, f"Error: {str(e)}"
+    result = captioner(image)
+    return result[0]['generated_text']
 
-with gr.Blocks(title="AI Image Modifier") as demo:
-    gr.Markdown("# 🎨 AI Image Generator & Modifier")
-    gr.Markdown("Create new images or modify existing ones using text prompts!")
-    
-    with gr.Tab("Generate New Image"):
-        with gr.Row():
-            with gr.Column():
-                gen_prompt = gr.Textbox(
-                    label="Describe your image",
-                    placeholder="A beautiful sunset over mountains, digital art...",
-                    lines=2
-                )
-                gen_negative = gr.Textbox(
-                    label="What to avoid (optional)",
-                    placeholder="blurry, low quality, distorted...",
-                    lines=1
-                )
-                with gr.Row():
-                    gen_steps = gr.Slider(1, 50, value=20, label="Steps")
-                    gen_guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale")
-                
-                gen_button = gr.Button("Generate Image", variant="primary")
-            
-            with gr.Column():
-                gen_output = gr.Image(label="Generated Image")
-                gen_status = gr.Textbox(label="Status", interactive=False)
-    
-    with gr.Tab("Modify Existing Image"):
-        with gr.Row():
-            with gr.Column():
-                mod_input_image = gr.Image(label="Upload Image", type="pil")
-                mod_prompt = gr.Textbox(
-                    label="How do you want to modify it?",
-                    placeholder="Make it look like winter, change style to oil painting...",
-                    lines=2
-                )
-                with gr.Row():
-                    mod_strength = gr.Slider(0.1, 1.0, value=0.75, label="Modification Strength")
-                    mod_steps = gr.Slider(1, 50, value=20, label="Steps")
-                    mod_guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale")
-                
-                mod_button = gr.Button("Modify Image", variant="primary")
-            
-            with gr.Column():
-                mod_output = gr.Image(label="Modified Image")
-                mod_status = gr.Textbox(label="Status", interactive=False)
-    
-    # Examples
-    gr.Markdown("### Example Prompts:")
-    gr.Examples(
-        examples=[
-            ["A majestic dragon flying over a medieval castle, fantasy art, highly detailed"],
-            ["A cyberpunk cityscape at night, neon lights, raining, futuristic"],
-            ["A cute corgi puppy wearing a superhero cape, cartoon style"],
-            ["An astronaut riding a horse on Mars, photorealistic"]
-        ],
-        inputs=gen_prompt
-    )
-    
-    # Connect functions
-    gen_button.click(
-        generate_image,
-        inputs=[gen_prompt, gen_negative, gen_steps, gen_guidance],
-        outputs=[gen_output, gen_status]
-    )
-    
-    mod_button.click(
-        modify_image,
-        inputs=[mod_input_image, mod_prompt, mod_strength, mod_steps, mod_guidance],
-        outputs=[mod_output, mod_status]
-    )
+demo = gr.Interface(
+    fn=caption_image,
+    inputs=gr.Image(label="Upload Image", type="pil"),
+    outputs=gr.Textbox(label="Generated Caption"),
+    title="Image Captioning",
+    description="Upload an image and AI will generate a caption for it",
+    examples=[
+        ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png"]
+    ]
+)
 
-demo.launch(share=True)
+demo.launch()