|
import os |
|
from huggingface_hub import login |
|
from transformers import BlipProcessor, BlipForConditionalGeneration |
|
from PIL import Image |
|
import pytesseract |
|
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
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 |
|
|
|
|
|
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium") |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
pipe.to(device) |
|
model.to(device) |
|
|
|
|
|
|
|
@spaces.GPU(duration=300) |
|
def generate_caption_and_image(image): |
|
img = image.convert("RGB") |
|
|
|
|
|
extracted_text = pytesseract.image_to_string(img) |
|
|
|
raw_image = image.convert("RGB") |
|
|
|
|
|
|
|
|
|
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) |
|
prompt = f"Create a highly realistic design of a clothing item based on the following description: 'The design should incorporate elements from the extracted text: {extracted_text}, along with the caption: {caption}. The clothing should look realistic, modern, and stylish. Use high-quality fabric textures and realistic lighting to give the design a lifelike appearance. The colors, patterns, and materials should reflect the essence of the caption and extracted text.'" |
|
|
|
|
|
|
|
generated_image = pipe(prompt).images[0] |
|
|
|
return caption, generated_image |
|
|
|
|
|
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) |
|
|
|
|