import gradio as gr
import requests
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
# from badnet_m import BadNet

# import timm

# model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
# model.train()

# model = BadNet(3, 10)

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
# pipeline = pipeline.to('cuda:0')


import os 



def print_bn():
    bn_data = []
    for m in model.modules():
        if(type(m) is nn.BatchNorm2d):
            # print(m.momentum)
            bn_data.extend(m.running_mean.data.numpy().tolist())
            bn_data.extend(m.running_var.data.numpy().tolist())
            bn_data.append(m.momentum)
    return bn_data

def greet(text):
    if(text == ''):
        # return 'changing'
        pipeline.unet.load_attn_procs('./pytorch_lora_weights.bin')
        pipeline.safety_checker = lambda images, clip_input: (images, False)
        # images = pipeline('a photo of dog').images
        # image = images[0]
        return None
    else:
        # return 'pipelining'
        images = pipeline(text).images
        image = images[0]
        return image

def greet_backdoor(image):
    # url = f'https://huggingface.co/spaces?p=1&sort=modified&search=GPT'
    # html = request_url(url)
    # key = os.getenv("OPENAI_API_KEY")
#     x = torch.ones([1,3,224,224])
    if(image is None):
        
        model.load_state_dict(torch.load("./badnet_cifar_all.pth"))
        return 'change to backdoor'
        
        # bn_data = print_bn()
        # return ','.join([f'{x:.10f}' for x in bn_data])
        
        
    else:  
        # print(type(image))
        # print(image.min(), image.max())
        # image = image[np.newaxis,:,:,:]
        # print(image.shape)
        # image = np.transpose(image,(0,3,1,2))
        
        
        image = torch.tensor(image).float()
        image = image/255.0
        image = image.unsqueeze(0)
        image = torch.permute(image, [0,3,1,2])

        
        
        
        
        # transform_nor = transforms.Compose([transforms.ToTensor(), transforms.Resize((32,32)), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])
        # image = transform_nor(image).unsqueeze(0)
        
        
        out = model(image)

    # model.train()
    return out



iface = gr.Interface(fn=greet, inputs='text', outputs="image")
# iface = gr.Interface(fn=greet, inputs='text', outputs="text")

# image = gr.inputs.Image(label="Upload a photo", shape=(32,32))
# iface = gr.Interface(fn=greet, inputs=image, outputs="text")
iface.launch()