File size: 1,571 Bytes
f667084
 
a04adbd
 
 
f667084
a04adbd
f667084
 
 
 
 
a04adbd
 
 
 
 
 
 
f667084
a04adbd
 
9023169
 
a04adbd
 
9023169
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a81e611
f667084
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import os
from huggingface_hub import login
from transformers import BlipProcessor, BlipForConditionalGeneration



# Get Hugging Face Token from environment variable
hf_token = os.getenv('HF_AUTH_TOKEN')
if not hf_token:
    raise ValueError("Hugging Face token is not set in the environment variables.")
login(token=hf_token)

# Load the processor and model
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
import gradio as gr
from diffusers import DiffusionPipeline
import torch
import spaces  # Hugging Face Spaces module

# Initialize the model
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")



@spaces.GPU(duration=300)
def generate_caption_and_image(image):
    # Process the image
    raw_image = image.convert("RGB")
    
    # Generate caption
    inputs = processor(raw_image, return_tensors="pt", padding=True, truncation=True, max_length=250)
    inputs = {key: val.to(device) for key, val in inputs.items()}
    out = model.generate(**inputs)
    caption = processor.decode(out[0], skip_special_tokens=True)

    # Generate image based on the caption
    generated_image = pipe(caption).images[0]

    return caption, generated_image

# Gradio UI
iface = gr.Interface(
    fn=generate_caption_and_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[gr.Textbox(label="Generated Caption"), gr.Image(label="Generated Design")],
    live=True
)
iface.launch(share=True)