StyleSync / app.py
Fiqa's picture
Update app.py
01b1364 verified
raw
history blame
3.17 kB
import os
from huggingface_hub import login
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline
import torch
import spaces # Hugging Face Spaces module
import requests
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from diffusers import DiffusionPipeline
fabrics = ['cotton', 'silk', 'denim', 'linen', 'polyester', 'wool', 'velvet']
patterns = ['striped', 'floral', 'geometric', 'abstract', 'solid', 'polka dots']
textile_designs = ['woven texture', 'embroidery', 'printed fabric', 'hand-dyed', 'quilting']
# 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")
processor1 = BlipProcessor.from_pretrained("noamrot/FuseCap")
model2 = BlipForConditionalGeneration.from_pretrained("noamrot/FuseCap")
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")
device = "cuda" if torch.cuda.is_available() else "cpu"
# pipe.to(device)
model2.to(device)
model.to(device)
pipe.to(device)
@spaces.GPU(duration=150)
def generate_caption_and_image(image):
img = image.convert("RGB")
# reader = easyocr.Reader(['en'])
# # result = reader.readtext(img)
import random
text = "a picture of "
inputs = processor(img, text, return_tensors="pt").to(device)
out = model2.generate(**inputs, num_beams = 3)
caption2 = processor.decode(out[0], skip_special_tokens=True)
# Generate caption
inputs = processor(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)
caption1 = processor.decode(out[0], skip_special_tokens=True)
prompt = f"Generate a clothing item using the following details: 1. {caption1} 2. {caption2} 3. Fabric: {selected_fabric} 4. Pattern: {selected_pattern} 5. Design Style: {selected_textile_design}. "
prompt +="The image should have a clean, minimalistic grey or white background, with realistic lighting and fine details, ensuring a sophisticated and polished appearance"
# Generate image based on the caption
generated_image = pipe(prompt).images[0]
return prompt, generated_image
# Gradio UI
iface = gr.Interface(
fn=generate_caption_and_image,
inputs=[gr.Image(type="pil", label="Upload Image"),
gr.Radio(fabrics, label="Select Fabric"),
gr.Radio(patterns, label="Select Pattern"),
gr.Radio(textile_designs, label="Select Textile Design")]
outputs=[gr.Textbox(label="Generated Caption"), gr.Image(label="Generated Design")],
live=True
)
iface.launch(share=True)